1use std::collections::{BTreeMap, BTreeSet, HashMap};
2
3use serde::{Deserialize, Serialize};
4
5use crate::pricing::{CostConfidence, UsdAmount};
6use crate::routing_explain::{RoutingExplainCandidate, RoutingExplainResponse};
7use crate::state::{
8 BalanceSnapshotStatus, FinishedRequest, ProviderBalanceSnapshot, UsageBucket, UsageRollupView,
9};
10use crate::usage_providers::UsageProviderRefreshSummary;
11
12#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default, Hash)]
13#[serde(rename_all = "snake_case")]
14pub enum UsageBalanceStatus {
15 #[default]
16 Unknown,
17 Ok,
18 Unlimited,
19 Exhausted,
20 Stale,
21 Error,
22}
23
24impl UsageBalanceStatus {
25 pub fn as_str(self) -> &'static str {
26 match self {
27 UsageBalanceStatus::Unknown => "unknown",
28 UsageBalanceStatus::Ok => "ok",
29 UsageBalanceStatus::Unlimited => "unlimited",
30 UsageBalanceStatus::Exhausted => "exhausted",
31 UsageBalanceStatus::Stale => "stale",
32 UsageBalanceStatus::Error => "error",
33 }
34 }
35
36 pub fn is_attention(self) -> bool {
37 matches!(
38 self,
39 UsageBalanceStatus::Unknown
40 | UsageBalanceStatus::Exhausted
41 | UsageBalanceStatus::Stale
42 | UsageBalanceStatus::Error
43 )
44 }
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
48pub struct UsageBalanceStatusCounts {
49 pub ok: usize,
50 pub unlimited: usize,
51 pub exhausted: usize,
52 pub stale: usize,
53 pub error: usize,
54 pub unknown: usize,
55}
56
57impl UsageBalanceStatusCounts {
58 pub fn total(&self) -> usize {
59 self.ok
60 .saturating_add(self.unlimited)
61 .saturating_add(self.exhausted)
62 .saturating_add(self.stale)
63 .saturating_add(self.error)
64 .saturating_add(self.unknown)
65 }
66
67 pub fn record(&mut self, status: UsageBalanceStatus) {
68 match status {
69 UsageBalanceStatus::Ok => self.ok += 1,
70 UsageBalanceStatus::Unlimited => self.unlimited += 1,
71 UsageBalanceStatus::Exhausted => self.exhausted += 1,
72 UsageBalanceStatus::Stale => self.stale += 1,
73 UsageBalanceStatus::Error => self.error += 1,
74 UsageBalanceStatus::Unknown => self.unknown += 1,
75 }
76 }
77
78 pub fn aggregate_status(&self) -> UsageBalanceStatus {
79 if self.total() == 0 {
80 UsageBalanceStatus::Unknown
81 } else if self.error > 0 {
82 UsageBalanceStatus::Error
83 } else if self.exhausted > 0 {
84 UsageBalanceStatus::Exhausted
85 } else if self.stale > 0 {
86 UsageBalanceStatus::Stale
87 } else if self.unknown > 0 {
88 UsageBalanceStatus::Unknown
89 } else if self.unlimited > 0 {
90 UsageBalanceStatus::Unlimited
91 } else {
92 UsageBalanceStatus::Ok
93 }
94 }
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
98pub struct UsageBalanceView {
99 pub service_name: String,
100 pub window_days: usize,
101 pub generated_at_ms: u64,
102 pub totals: UsageBalanceTotals,
103 pub provider_rows: Vec<UsageBalanceProviderRow>,
104 pub endpoint_rows: Vec<UsageBalanceEndpointRow>,
105 pub routing_impacts: Vec<UsageBalanceRouteImpact>,
106 pub refresh_status: UsageBalanceRefreshStatus,
107 #[serde(default, skip_serializing_if = "Vec::is_empty")]
108 pub warnings: Vec<String>,
109}
110
111#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
112pub struct UsageBalanceTotals {
113 pub requests_total: u64,
114 pub requests_error: u64,
115 pub success_per_mille: Option<u16>,
116 pub input_tokens: i64,
117 pub cached_input_tokens: i64,
118 pub output_tokens: i64,
119 pub reasoning_output_tokens: i64,
120 pub total_tokens: i64,
121 pub cost_total_usd: Option<String>,
122 pub cost_display: String,
123 pub cost_confidence: CostConfidence,
124 pub balance_status_counts: UsageBalanceStatusCounts,
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
128pub struct UsageBalanceProviderRow {
129 pub provider_id: String,
130 pub display_name: String,
131 pub usage: UsageBucket,
132 pub success_per_mille: Option<u16>,
133 pub output_tokens_per_second: Option<u64>,
134 pub avg_ttfb_ms: Option<u64>,
135 pub cost_total_usd: Option<String>,
136 pub cost_display: String,
137 pub cost_confidence: CostConfidence,
138 pub balance_status: UsageBalanceStatus,
139 pub balance_counts: UsageBalanceStatusCounts,
140 pub primary_balance: Option<UsageBalanceSnapshotSummary>,
141 pub latest_balance_error: Option<String>,
142 pub balance_age_ms: Option<u64>,
143 pub routing: UsageBalanceRouteImpact,
144 pub endpoint_count: usize,
145 pub endpoints_with_balance: usize,
146 pub recent_endpoint_requests: u64,
147}
148
149impl UsageBalanceProviderRow {
150 pub fn needs_attention(&self) -> bool {
151 self.balance_status.is_attention()
152 || self
153 .latest_balance_error
154 .as_deref()
155 .is_some_and(|value| !value.trim().is_empty())
156 || self.usage.requests_error > 0
157 }
158}
159
160#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
161pub struct UsageBalanceEndpointRow {
162 pub provider_id: String,
163 pub endpoint_id: String,
164 pub station_name: Option<String>,
165 pub upstream_index: Option<usize>,
166 pub base_url: Option<String>,
167 pub usage: UsageBucket,
168 pub balance_status: UsageBalanceStatus,
169 pub balance: Option<UsageBalanceSnapshotSummary>,
170 pub route_selected: bool,
171 pub route_skip_reasons: Vec<String>,
172}
173
174#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
175pub struct UsageBalanceSnapshotSummary {
176 pub provider_id: String,
177 pub station_name: Option<String>,
178 pub upstream_index: Option<usize>,
179 pub source: String,
180 pub status: UsageBalanceStatus,
181 pub amount_summary: String,
182 pub fetched_at_ms: u64,
183 pub stale_after_ms: Option<u64>,
184 pub age_ms: Option<u64>,
185 pub error: Option<String>,
186 pub exhaustion_affects_routing: bool,
187 pub routing_exhausted: bool,
188 pub routing_ignored_exhaustion: bool,
189}
190
191#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
192pub struct UsageBalanceRouteImpact {
193 pub provider_id: String,
194 pub selected: bool,
195 pub selected_endpoint_id: Option<String>,
196 pub selected_provider_endpoint_key: Option<String>,
197 pub candidate_count: usize,
198 pub skip_reasons: Vec<String>,
199 pub route_paths: Vec<Vec<String>>,
200}
201
202#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
203pub struct UsageBalanceRefreshStatus {
204 pub refreshing: bool,
205 pub total_snapshots: usize,
206 pub latest_fetched_at_ms: Option<u64>,
207 pub latest_error: Option<String>,
208 #[serde(default, skip_serializing_if = "Option::is_none")]
209 pub latest_error_provider_id: Option<String>,
210 pub last_message: Option<String>,
211 pub last_error: Option<String>,
212 #[serde(default, skip_serializing_if = "Option::is_none")]
213 pub last_provider_refresh: Option<UsageProviderRefreshSummary>,
214 pub status_counts: UsageBalanceStatusCounts,
215}
216
217#[derive(Debug, Clone, Default)]
218pub struct UsageBalanceRefreshInput {
219 pub refreshing: bool,
220 pub last_message: Option<String>,
221 pub last_error: Option<String>,
222 pub last_provider_refresh: Option<UsageProviderRefreshSummary>,
223}
224
225pub struct UsageBalanceBuildInput<'a> {
226 pub service_name: &'a str,
227 pub window_days: usize,
228 pub generated_at_ms: u64,
229 pub usage_rollup: &'a UsageRollupView,
230 pub provider_balances: &'a HashMap<String, Vec<ProviderBalanceSnapshot>>,
231 pub recent: &'a [FinishedRequest],
232 pub routing_explain: Option<&'a RoutingExplainResponse>,
233 pub refresh: UsageBalanceRefreshInput,
234}
235
236impl UsageBalanceView {
237 pub fn build(input: UsageBalanceBuildInput<'_>) -> Self {
238 let route_index = RouteIndex::from_explain(input.routing_explain);
239 let mut endpoint_rows = build_endpoint_rows(
240 input.recent,
241 input.provider_balances,
242 &route_index,
243 input.generated_at_ms,
244 );
245 sort_endpoint_rows(&mut endpoint_rows);
246
247 let provider_rows = build_provider_rows(&input, &route_index, &endpoint_rows);
248 let totals = build_totals(
249 input.usage_rollup,
250 input.provider_balances,
251 input.generated_at_ms,
252 );
253 let refresh_status = build_refresh_status(
254 input.provider_balances,
255 input.generated_at_ms,
256 input.refresh,
257 );
258 let routing_impacts = route_index
259 .provider_impacts
260 .values()
261 .cloned()
262 .collect::<Vec<_>>();
263
264 Self {
265 service_name: input.service_name.to_string(),
266 window_days: input.window_days,
267 generated_at_ms: input.generated_at_ms,
268 totals,
269 provider_rows,
270 endpoint_rows,
271 routing_impacts,
272 refresh_status,
273 warnings: Vec::new(),
274 }
275 }
276}
277
278fn build_totals(
279 rollup: &UsageRollupView,
280 provider_balances: &HashMap<String, Vec<ProviderBalanceSnapshot>>,
281 now_ms: u64,
282) -> UsageBalanceTotals {
283 let window = &rollup.window;
284 let mut balance_status_counts = UsageBalanceStatusCounts::default();
285 for snapshot in provider_balances.values().flatten() {
286 balance_status_counts.record(classify_snapshot(snapshot, now_ms));
287 }
288
289 UsageBalanceTotals {
290 requests_total: window.requests_total,
291 requests_error: window.requests_error,
292 success_per_mille: per_mille(
293 window.requests_total.saturating_sub(window.requests_error),
294 window.requests_total,
295 ),
296 input_tokens: window.usage.input_tokens,
297 cached_input_tokens: window.usage.cache_read_tokens_total(),
298 output_tokens: window.usage.output_tokens,
299 reasoning_output_tokens: window.usage.reasoning_output_tokens_total(),
300 total_tokens: window.usage.total_tokens,
301 cost_total_usd: window.cost.total_cost_usd.clone(),
302 cost_display: window.cost.display_total(),
303 cost_confidence: window.cost.confidence,
304 balance_status_counts,
305 }
306}
307
308fn build_provider_rows(
309 input: &UsageBalanceBuildInput<'_>,
310 route_index: &RouteIndex,
311 endpoint_rows: &[UsageBalanceEndpointRow],
312) -> Vec<UsageBalanceProviderRow> {
313 let mut provider_ids = BTreeSet::new();
314 for (provider_id, _) in &input.usage_rollup.by_provider {
315 provider_ids.insert(provider_id.clone());
316 }
317 for (provider_id, _) in group_balances_by_provider(input.provider_balances) {
318 provider_ids.insert(provider_id);
319 }
320 for request in input.recent {
321 if let Some(provider_id) = request
322 .provider_id
323 .as_deref()
324 .map(str::trim)
325 .filter(|value| !value.is_empty())
326 {
327 provider_ids.insert(provider_id.to_string());
328 }
329 }
330 for provider_id in route_index.provider_impacts.keys() {
331 provider_ids.insert(provider_id.clone());
332 }
333
334 let usage_by_provider = input
335 .usage_rollup
336 .by_provider
337 .iter()
338 .cloned()
339 .collect::<BTreeMap<_, _>>();
340 let balances_by_provider = group_balances_by_provider(input.provider_balances);
341 let mut endpoint_counts = BTreeMap::<String, (usize, usize, u64)>::new();
342 for endpoint in endpoint_rows {
343 let entry = endpoint_counts
344 .entry(endpoint.provider_id.clone())
345 .or_insert((0, 0, 0));
346 entry.0 += 1;
347 if endpoint.balance.is_some() {
348 entry.1 += 1;
349 }
350 entry.2 = entry.2.saturating_add(endpoint.usage.requests_total);
351 }
352
353 let mut rows = provider_ids
354 .into_iter()
355 .map(|provider_id| {
356 let usage = usage_by_provider
357 .get(provider_id.as_str())
358 .cloned()
359 .unwrap_or_default();
360 let balances = balances_by_provider
361 .get(provider_id.as_str())
362 .cloned()
363 .unwrap_or_default();
364 let balance_counts = balance_counts_for_snapshots(&balances, input.generated_at_ms);
365 let primary_balance = primary_balance_snapshot(&balances, input.generated_at_ms)
366 .map(|snapshot| summarize_snapshot(snapshot, input.generated_at_ms));
367 let latest_balance_error = latest_balance_error(&balances);
368 let balance_age_ms = primary_balance.as_ref().and_then(|summary| summary.age_ms);
369 let balance_status = balance_counts.aggregate_status();
370 let routing = route_index
371 .provider_impacts
372 .get(provider_id.as_str())
373 .cloned()
374 .unwrap_or_else(|| UsageBalanceRouteImpact {
375 provider_id: provider_id.clone(),
376 ..UsageBalanceRouteImpact::default()
377 });
378 let (endpoint_count, endpoints_with_balance, recent_endpoint_requests) =
379 endpoint_counts
380 .get(provider_id.as_str())
381 .copied()
382 .unwrap_or_default();
383
384 UsageBalanceProviderRow {
385 display_name: provider_id.clone(),
386 provider_id,
387 success_per_mille: per_mille(
388 usage.requests_total.saturating_sub(usage.requests_error),
389 usage.requests_total,
390 ),
391 output_tokens_per_second: output_tokens_per_second(&usage),
392 avg_ttfb_ms: (usage.ttfb_samples > 0)
393 .then(|| usage.ttfb_ms_total / usage.ttfb_samples),
394 cost_total_usd: usage.cost.total_cost_usd.clone(),
395 cost_display: usage.cost.display_total(),
396 cost_confidence: usage.cost.confidence,
397 usage,
398 balance_status,
399 balance_counts,
400 primary_balance,
401 latest_balance_error,
402 balance_age_ms,
403 routing,
404 endpoint_count,
405 endpoints_with_balance,
406 recent_endpoint_requests,
407 }
408 })
409 .collect::<Vec<_>>();
410
411 sort_provider_rows(&mut rows);
412 rows
413}
414
415fn build_refresh_status(
416 provider_balances: &HashMap<String, Vec<ProviderBalanceSnapshot>>,
417 now_ms: u64,
418 input: UsageBalanceRefreshInput,
419) -> UsageBalanceRefreshStatus {
420 let mut out = UsageBalanceRefreshStatus {
421 refreshing: input.refreshing,
422 last_message: input.last_message,
423 last_error: input.last_error,
424 last_provider_refresh: input.last_provider_refresh,
425 ..UsageBalanceRefreshStatus::default()
426 };
427 let mut latest_error_fetched_at_ms: Option<u64> = None;
428
429 for snapshot in provider_balances.values().flatten() {
430 out.total_snapshots += 1;
431 out.latest_fetched_at_ms = Some(
432 out.latest_fetched_at_ms
433 .unwrap_or(0)
434 .max(snapshot.fetched_at_ms),
435 );
436 let status = classify_snapshot(snapshot, now_ms);
437 out.status_counts.record(status);
438 if let Some(error) = snapshot
439 .error
440 .as_deref()
441 .map(str::trim)
442 .filter(|v| !v.is_empty())
443 && latest_error_fetched_at_ms.is_none_or(|latest| snapshot.fetched_at_ms >= latest)
444 {
445 out.latest_error = Some(error.to_string());
446 out.latest_error_provider_id = non_empty(snapshot.provider_id.clone())
447 .or_else(|| non_empty(snapshot.station_name.clone().unwrap_or_default()));
448 latest_error_fetched_at_ms = Some(snapshot.fetched_at_ms);
449 }
450 }
451
452 out
453}
454
455fn build_endpoint_rows(
456 recent: &[FinishedRequest],
457 provider_balances: &HashMap<String, Vec<ProviderBalanceSnapshot>>,
458 route_index: &RouteIndex,
459 now_ms: u64,
460) -> Vec<UsageBalanceEndpointRow> {
461 let mut acc = BTreeMap::<EndpointKey, EndpointAccum>::new();
462
463 for candidate in route_index.candidates.values() {
464 let key = EndpointKey {
465 provider_id: candidate.provider_id.clone(),
466 endpoint_id: candidate.endpoint_id.clone(),
467 };
468 let row = acc
469 .entry(key.clone())
470 .or_insert_with(|| EndpointAccum::new(key));
471 row.base_url = non_empty(candidate.base_url.clone()).or(row.base_url.take());
472 row.route_selected |= candidate.selected;
473 row.route_skip_reasons
474 .extend(candidate.skip_reasons.iter().cloned());
475 if let Some((station_name, upstream_index)) = candidate.compatibility.as_ref() {
476 row.station_name = Some(station_name.clone());
477 row.upstream_index = Some(*upstream_index);
478 }
479 }
480
481 for request in recent {
482 let Some(provider_id) = request
483 .provider_id
484 .as_deref()
485 .map(str::trim)
486 .filter(|value| !value.is_empty())
487 else {
488 continue;
489 };
490 let base_url = request
491 .upstream_base_url
492 .as_deref()
493 .map(str::trim)
494 .filter(|value| !value.is_empty());
495 let endpoint_id = base_url
496 .and_then(|base_url| {
497 route_index
498 .endpoint_by_provider_base_url
499 .get(&(provider_id.to_string(), base_url.to_string()))
500 .cloned()
501 })
502 .unwrap_or_else(|| {
503 base_url
504 .map(endpoint_id_from_base_url)
505 .unwrap_or_else(|| "unknown_endpoint".to_string())
506 });
507 let key = EndpointKey {
508 provider_id: provider_id.to_string(),
509 endpoint_id,
510 };
511 let row = acc
512 .entry(key.clone())
513 .or_insert_with(|| EndpointAccum::new(key));
514 if row.base_url.is_none() {
515 row.base_url = base_url.map(ToOwned::to_owned);
516 }
517 record_finished_into_bucket(&mut row.usage, request);
518 }
519
520 for (station_name, snapshots) in provider_balances {
521 for snapshot in snapshots {
522 let provider_id = balance_provider_id(station_name, snapshot);
523 let endpoint_id = snapshot
524 .upstream_index
525 .and_then(|idx| {
526 route_index
527 .endpoint_by_provider_station_upstream
528 .get(&(provider_id.clone(), station_name.clone(), idx))
529 .cloned()
530 })
531 .unwrap_or_else(|| {
532 snapshot
533 .upstream_index
534 .map(|idx| format!("upstream#{idx}"))
535 .or_else(|| non_empty(snapshot.source.clone()))
536 .unwrap_or_else(|| "balance".to_string())
537 });
538 let key = EndpointKey {
539 provider_id: provider_id.clone(),
540 endpoint_id,
541 };
542 let row = acc
543 .entry(key.clone())
544 .or_insert_with(|| EndpointAccum::new(key));
545 row.station_name = snapshot
546 .station_name
547 .clone()
548 .or_else(|| Some(station_name.clone()))
549 .or(row.station_name.take());
550 row.upstream_index = snapshot.upstream_index.or(row.upstream_index);
551 let summary = summarize_snapshot(snapshot, now_ms);
552 let replace = row.balance.as_ref().is_none_or(|existing| {
553 snapshot_display_rank(summary.status) < snapshot_display_rank(existing.status)
554 });
555 if replace {
556 row.balance = Some(summary);
557 }
558 }
559 }
560
561 acc.into_values().map(EndpointAccum::finish).collect()
562}
563
564fn record_finished_into_bucket(bucket: &mut UsageBucket, request: &FinishedRequest) {
565 bucket.requests_total = bucket.requests_total.saturating_add(1);
566 if request.status_code >= 400 {
567 bucket.requests_error = bucket.requests_error.saturating_add(1);
568 }
569 bucket.duration_ms_total = bucket.duration_ms_total.saturating_add(request.duration_ms);
570
571 let Some(usage) = request.usage.as_ref() else {
572 return;
573 };
574
575 bucket.usage.add_assign(usage);
576 bucket.cost.record_usage_cost(&request.cost);
577 bucket.requests_with_usage = bucket.requests_with_usage.saturating_add(1);
578 bucket.duration_ms_with_usage_total = bucket
579 .duration_ms_with_usage_total
580 .saturating_add(request.duration_ms);
581 let observability = request.observability_view();
582 let generation_ms = observability.generation_ms.unwrap_or(request.duration_ms);
583 bucket.generation_ms_total = bucket.generation_ms_total.saturating_add(generation_ms);
584 if let Some(ttfb) = observability.ttfb_ms.filter(|value| *value > 0) {
585 bucket.ttfb_ms_total = bucket.ttfb_ms_total.saturating_add(ttfb);
586 bucket.ttfb_samples = bucket.ttfb_samples.saturating_add(1);
587 }
588}
589
590fn group_balances_by_provider(
591 provider_balances: &HashMap<String, Vec<ProviderBalanceSnapshot>>,
592) -> BTreeMap<String, Vec<&ProviderBalanceSnapshot>> {
593 let mut out = BTreeMap::<String, Vec<&ProviderBalanceSnapshot>>::new();
594 for (station_name, snapshots) in provider_balances {
595 for snapshot in snapshots {
596 out.entry(balance_provider_id(station_name, snapshot))
597 .or_default()
598 .push(snapshot);
599 }
600 }
601 out
602}
603
604fn balance_provider_id(station_name: &str, snapshot: &ProviderBalanceSnapshot) -> String {
605 if snapshot.provider_id.trim().is_empty() {
606 station_name.to_string()
607 } else {
608 snapshot.provider_id.trim().to_string()
609 }
610}
611
612fn balance_counts_for_snapshots(
613 snapshots: &[&ProviderBalanceSnapshot],
614 now_ms: u64,
615) -> UsageBalanceStatusCounts {
616 let mut counts = UsageBalanceStatusCounts::default();
617 for snapshot in snapshots {
618 counts.record(classify_snapshot(snapshot, now_ms));
619 }
620 counts
621}
622
623fn classify_snapshot(snapshot: &ProviderBalanceSnapshot, now_ms: u64) -> UsageBalanceStatus {
624 match snapshot.status_at(now_ms) {
625 BalanceSnapshotStatus::Ok if snapshot.unlimited_quota == Some(true) => {
626 UsageBalanceStatus::Unlimited
627 }
628 BalanceSnapshotStatus::Ok => UsageBalanceStatus::Ok,
629 BalanceSnapshotStatus::Exhausted => UsageBalanceStatus::Exhausted,
630 BalanceSnapshotStatus::Stale => UsageBalanceStatus::Stale,
631 BalanceSnapshotStatus::Error => UsageBalanceStatus::Error,
632 BalanceSnapshotStatus::Unknown => UsageBalanceStatus::Unknown,
633 }
634}
635
636fn summarize_snapshot(
637 snapshot: &ProviderBalanceSnapshot,
638 now_ms: u64,
639) -> UsageBalanceSnapshotSummary {
640 UsageBalanceSnapshotSummary {
641 provider_id: snapshot.provider_id.clone(),
642 station_name: snapshot.station_name.clone(),
643 upstream_index: snapshot.upstream_index,
644 source: snapshot.source.clone(),
645 status: classify_snapshot(snapshot, now_ms),
646 amount_summary: snapshot.amount_summary(),
647 fetched_at_ms: snapshot.fetched_at_ms,
648 stale_after_ms: snapshot.stale_after_ms,
649 age_ms: (now_ms >= snapshot.fetched_at_ms).then(|| now_ms - snapshot.fetched_at_ms),
650 error: snapshot.error.clone(),
651 exhaustion_affects_routing: snapshot.exhaustion_affects_routing,
652 routing_exhausted: snapshot.exhaustion_affects_routing
653 && snapshot.status_at(now_ms) == BalanceSnapshotStatus::Exhausted,
654 routing_ignored_exhaustion: !snapshot.exhaustion_affects_routing
655 && snapshot.status_at(now_ms) == BalanceSnapshotStatus::Exhausted,
656 }
657}
658
659fn primary_balance_snapshot<'a>(
660 snapshots: &'a [&ProviderBalanceSnapshot],
661 now_ms: u64,
662) -> Option<&'a ProviderBalanceSnapshot> {
663 snapshots.iter().copied().min_by(|left, right| {
664 snapshot_display_rank(classify_snapshot(left, now_ms))
665 .cmp(&snapshot_display_rank(classify_snapshot(right, now_ms)))
666 .then_with(|| left.upstream_index.cmp(&right.upstream_index))
667 .then_with(|| right.fetched_at_ms.cmp(&left.fetched_at_ms))
668 })
669}
670
671fn snapshot_display_rank(status: UsageBalanceStatus) -> u8 {
672 match status {
673 UsageBalanceStatus::Ok | UsageBalanceStatus::Unlimited => 0,
674 UsageBalanceStatus::Stale => 1,
675 UsageBalanceStatus::Unknown | UsageBalanceStatus::Error => 2,
676 UsageBalanceStatus::Exhausted => 3,
677 }
678}
679
680fn latest_balance_error(snapshots: &[&ProviderBalanceSnapshot]) -> Option<String> {
681 snapshots
682 .iter()
683 .filter_map(|snapshot| {
684 snapshot
685 .error
686 .as_deref()
687 .map(str::trim)
688 .filter(|value| !value.is_empty())
689 .map(|error| (snapshot.fetched_at_ms, error.to_string()))
690 })
691 .max_by_key(|(fetched_at_ms, _)| *fetched_at_ms)
692 .map(|(_, error)| error)
693}
694
695fn per_mille(num: u64, den: u64) -> Option<u16> {
696 (den > 0).then(|| ((num.saturating_mul(1000) / den).min(1000)) as u16)
697}
698
699fn output_tokens_per_second(bucket: &UsageBucket) -> Option<u64> {
700 let output = bucket.usage.output_tokens.max(0) as u64;
701 if output == 0 || bucket.generation_ms_total == 0 {
702 return None;
703 }
704 Some(output.saturating_mul(1000) / bucket.generation_ms_total)
705}
706
707fn sort_provider_rows(rows: &mut [UsageBalanceProviderRow]) {
708 rows.sort_by(|left, right| {
709 provider_selected_rank(right)
710 .cmp(&provider_selected_rank(left))
711 .then_with(|| {
712 provider_attention_rank(left.balance_status)
713 .cmp(&provider_attention_rank(right.balance_status))
714 })
715 .then_with(|| provider_cost_sort_value(right).cmp(&provider_cost_sort_value(left)))
716 .then_with(|| right.usage.requests_total.cmp(&left.usage.requests_total))
717 .then_with(|| left.provider_id.cmp(&right.provider_id))
718 });
719}
720
721fn provider_selected_rank(row: &UsageBalanceProviderRow) -> u8 {
722 u8::from(row.routing.selected)
723}
724
725fn provider_attention_rank(status: UsageBalanceStatus) -> u8 {
726 match status {
727 UsageBalanceStatus::Error => 0,
728 UsageBalanceStatus::Exhausted => 1,
729 UsageBalanceStatus::Stale => 2,
730 UsageBalanceStatus::Unknown => 3,
731 UsageBalanceStatus::Unlimited => 4,
732 UsageBalanceStatus::Ok => 5,
733 }
734}
735
736fn provider_cost_sort_value(row: &UsageBalanceProviderRow) -> i128 {
737 row.cost_total_usd
738 .as_deref()
739 .and_then(UsdAmount::from_decimal_str)
740 .map(UsdAmount::femto_usd)
741 .unwrap_or(0)
742}
743
744fn sort_endpoint_rows(rows: &mut [UsageBalanceEndpointRow]) {
745 rows.sort_by(|left, right| {
746 u8::from(right.route_selected)
747 .cmp(&u8::from(left.route_selected))
748 .then_with(|| {
749 provider_attention_rank(left.balance_status)
750 .cmp(&provider_attention_rank(right.balance_status))
751 })
752 .then_with(|| right.usage.requests_total.cmp(&left.usage.requests_total))
753 .then_with(|| left.provider_id.cmp(&right.provider_id))
754 .then_with(|| left.endpoint_id.cmp(&right.endpoint_id))
755 });
756}
757
758#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
759struct EndpointKey {
760 provider_id: String,
761 endpoint_id: String,
762}
763
764#[derive(Debug, Clone)]
765struct EndpointAccum {
766 key: EndpointKey,
767 station_name: Option<String>,
768 upstream_index: Option<usize>,
769 base_url: Option<String>,
770 usage: UsageBucket,
771 balance: Option<UsageBalanceSnapshotSummary>,
772 route_selected: bool,
773 route_skip_reasons: BTreeSet<String>,
774}
775
776impl EndpointAccum {
777 fn new(key: EndpointKey) -> Self {
778 Self {
779 key,
780 station_name: None,
781 upstream_index: None,
782 base_url: None,
783 usage: UsageBucket::default(),
784 balance: None,
785 route_selected: false,
786 route_skip_reasons: BTreeSet::new(),
787 }
788 }
789
790 fn finish(self) -> UsageBalanceEndpointRow {
791 let balance_status = self
792 .balance
793 .as_ref()
794 .map(|balance| balance.status)
795 .unwrap_or(UsageBalanceStatus::Unknown);
796 UsageBalanceEndpointRow {
797 provider_id: self.key.provider_id,
798 endpoint_id: self.key.endpoint_id,
799 station_name: self.station_name,
800 upstream_index: self.upstream_index,
801 base_url: self.base_url,
802 usage: self.usage,
803 balance_status,
804 balance: self.balance,
805 route_selected: self.route_selected,
806 route_skip_reasons: self.route_skip_reasons.into_iter().collect(),
807 }
808 }
809}
810
811#[derive(Debug, Clone)]
812struct RouteCandidateView {
813 provider_id: String,
814 endpoint_id: String,
815 base_url: String,
816 selected: bool,
817 skip_reasons: Vec<String>,
818 compatibility: Option<(String, usize)>,
819}
820
821#[derive(Debug, Clone, Default)]
822struct RouteIndex {
823 candidates: BTreeMap<String, RouteCandidateView>,
824 provider_impacts: BTreeMap<String, UsageBalanceRouteImpact>,
825 endpoint_by_provider_base_url: BTreeMap<(String, String), String>,
826 endpoint_by_provider_station_upstream: BTreeMap<(String, String, usize), String>,
827}
828
829impl RouteIndex {
830 fn from_explain(explain: Option<&RoutingExplainResponse>) -> Self {
831 let mut out = Self::default();
832 let Some(explain) = explain else {
833 return out;
834 };
835
836 for candidate in &explain.candidates {
837 out.record_candidate(candidate);
838 }
839 for impact in out.provider_impacts.values_mut() {
840 impact.skip_reasons.sort();
841 impact.skip_reasons.dedup();
842 impact.route_paths.sort();
843 impact.route_paths.dedup();
844 }
845 out
846 }
847
848 fn record_candidate(&mut self, candidate: &RoutingExplainCandidate) {
849 let skip_reasons = candidate
850 .skip_reasons
851 .iter()
852 .map(|reason| reason.code().to_string())
853 .collect::<Vec<_>>();
854 let compatibility = candidate.compatibility.as_ref().map(|compatibility| {
855 (
856 compatibility.station_name.clone(),
857 compatibility.upstream_index,
858 )
859 });
860 let view = RouteCandidateView {
861 provider_id: candidate.provider_id.clone(),
862 endpoint_id: candidate.endpoint_id.clone(),
863 base_url: candidate.upstream_base_url.clone(),
864 selected: candidate.selected,
865 skip_reasons: skip_reasons.clone(),
866 compatibility: compatibility.clone(),
867 };
868
869 self.endpoint_by_provider_base_url.insert(
870 (
871 candidate.provider_id.clone(),
872 candidate.upstream_base_url.clone(),
873 ),
874 candidate.endpoint_id.clone(),
875 );
876 if let Some((station_name, upstream_index)) = compatibility {
877 self.endpoint_by_provider_station_upstream.insert(
878 (candidate.provider_id.clone(), station_name, upstream_index),
879 candidate.endpoint_id.clone(),
880 );
881 }
882
883 let impact = self
884 .provider_impacts
885 .entry(candidate.provider_id.clone())
886 .or_insert_with(|| UsageBalanceRouteImpact {
887 provider_id: candidate.provider_id.clone(),
888 ..UsageBalanceRouteImpact::default()
889 });
890 impact.candidate_count += 1;
891 impact.route_paths.push(candidate.route_path.clone());
892 impact.skip_reasons.extend(skip_reasons);
893 if candidate.selected {
894 impact.selected = true;
895 impact.selected_endpoint_id = Some(candidate.endpoint_id.clone());
896 impact.selected_provider_endpoint_key = Some(candidate.provider_endpoint_key.clone());
897 }
898
899 self.candidates
900 .insert(candidate.provider_endpoint_key.clone(), view);
901 }
902}
903
904fn endpoint_id_from_base_url(base_url: &str) -> String {
905 let trimmed = base_url.trim();
906 let after_scheme = trimmed
907 .split_once("://")
908 .map(|(_, rest)| rest)
909 .unwrap_or(trimmed);
910 after_scheme
911 .split('/')
912 .next()
913 .map(str::trim)
914 .filter(|value| !value.is_empty())
915 .unwrap_or("endpoint")
916 .to_string()
917}
918
919fn non_empty(value: String) -> Option<String> {
920 (!value.trim().is_empty()).then_some(value)
921}
922
923#[cfg(test)]
924mod tests {
925 use super::*;
926 use crate::pricing::CostBreakdown;
927 use crate::routing_explain::{
928 RoutingExplainCandidate, RoutingExplainResponse, RoutingExplainSkipReason,
929 };
930 use crate::usage::UsageMetrics;
931
932 fn build_view(
933 usage_rollup: &UsageRollupView,
934 provider_balances: &HashMap<String, Vec<ProviderBalanceSnapshot>>,
935 recent: &[FinishedRequest],
936 routing_explain: Option<&RoutingExplainResponse>,
937 ) -> UsageBalanceView {
938 UsageBalanceView::build(UsageBalanceBuildInput {
939 service_name: "codex",
940 window_days: 7,
941 generated_at_ms: 1_000,
942 usage_rollup,
943 provider_balances,
944 recent,
945 routing_explain,
946 refresh: UsageBalanceRefreshInput::default(),
947 })
948 }
949
950 #[test]
951 fn balance_status_counts_keep_unknown_stale_exhausted_error_and_unlimited_distinct() {
952 let mut balances = HashMap::new();
953 balances.insert(
954 "station".to_string(),
955 vec![
956 ProviderBalanceSnapshot {
957 provider_id: "ok".to_string(),
958 exhausted: Some(false),
959 total_balance_usd: Some("3".to_string()),
960 fetched_at_ms: 900,
961 ..ProviderBalanceSnapshot::default()
962 },
963 ProviderBalanceSnapshot {
964 provider_id: "unlimited".to_string(),
965 exhausted: Some(false),
966 unlimited_quota: Some(true),
967 fetched_at_ms: 900,
968 ..ProviderBalanceSnapshot::default()
969 },
970 ProviderBalanceSnapshot {
971 provider_id: "stale".to_string(),
972 exhausted: Some(false),
973 stale_after_ms: Some(950),
974 fetched_at_ms: 900,
975 ..ProviderBalanceSnapshot::default()
976 },
977 ProviderBalanceSnapshot {
978 provider_id: "exhausted".to_string(),
979 exhausted: Some(true),
980 fetched_at_ms: 900,
981 ..ProviderBalanceSnapshot::default()
982 },
983 ProviderBalanceSnapshot {
984 provider_id: "error".to_string(),
985 error: Some("lookup failed".to_string()),
986 fetched_at_ms: 900,
987 ..ProviderBalanceSnapshot::default()
988 },
989 ProviderBalanceSnapshot {
990 provider_id: "unknown".to_string(),
991 fetched_at_ms: 900,
992 ..ProviderBalanceSnapshot::default()
993 },
994 ],
995 );
996
997 let view = build_view(&UsageRollupView::default(), &balances, &[], None);
998
999 assert_eq!(view.totals.balance_status_counts.ok, 1);
1000 assert_eq!(view.totals.balance_status_counts.unlimited, 1);
1001 assert_eq!(view.totals.balance_status_counts.stale, 1);
1002 assert_eq!(view.totals.balance_status_counts.exhausted, 1);
1003 assert_eq!(view.totals.balance_status_counts.error, 1);
1004 assert_eq!(view.totals.balance_status_counts.unknown, 1);
1005 assert_eq!(
1006 view.provider_rows
1007 .iter()
1008 .find(|row| row.provider_id == "unlimited")
1009 .map(|row| row.balance_status),
1010 Some(UsageBalanceStatus::Unlimited)
1011 );
1012 assert_eq!(
1013 view.provider_rows
1014 .iter()
1015 .find(|row| row.provider_id == "unknown")
1016 .map(|row| row.balance_status),
1017 Some(UsageBalanceStatus::Unknown)
1018 );
1019 }
1020
1021 #[test]
1022 fn provider_rows_prefer_route_selection_then_attention_and_usage() {
1023 let rollup = UsageRollupView {
1024 by_provider: vec![
1025 (
1026 "cheap".to_string(),
1027 UsageBucket {
1028 requests_total: 2,
1029 ..UsageBucket::default()
1030 },
1031 ),
1032 (
1033 "selected".to_string(),
1034 UsageBucket {
1035 requests_total: 1,
1036 ..UsageBucket::default()
1037 },
1038 ),
1039 ],
1040 ..UsageRollupView::default()
1041 };
1042 let explain = RoutingExplainResponse {
1043 api_version: 1,
1044 service_name: "codex".to_string(),
1045 runtime_loaded_at_ms: None,
1046 request_model: None,
1047 session_id: None,
1048 request_context: Default::default(),
1049 selected_route: None,
1050 candidates: vec![RoutingExplainCandidate {
1051 provider_id: "selected".to_string(),
1052 provider_alias: None,
1053 endpoint_id: "default".to_string(),
1054 provider_endpoint_key: "codex:selected:default".to_string(),
1055 route_path: vec!["main".to_string(), "selected".to_string()],
1056 preference_group: 0,
1057 compatibility: None,
1058 upstream_base_url: "https://selected.example/v1".to_string(),
1059 selected: true,
1060 skip_reasons: Vec::new(),
1061 }],
1062 affinity_policy: "off".to_string(),
1063 affinity: None,
1064 conditional_routes: Vec::new(),
1065 };
1066
1067 let view = build_view(&rollup, &HashMap::new(), &[], Some(&explain));
1068
1069 assert_eq!(view.provider_rows[0].provider_id, "selected");
1070 assert!(view.provider_rows[0].routing.selected);
1071 }
1072
1073 #[test]
1074 fn endpoint_rows_merge_recent_usage_balance_and_route_skip_reasons() {
1075 let recent = vec![FinishedRequest {
1076 id: 1,
1077 trace_id: None,
1078 session_id: None,
1079 session_identity_source: None,
1080 client_name: None,
1081 client_addr: None,
1082 cwd: None,
1083 model: Some("gpt-test".to_string()),
1084 reasoning_effort: None,
1085 service_tier: None,
1086 station_name: Some("station-a".to_string()),
1087 provider_id: Some("right".to_string()),
1088 upstream_base_url: Some("https://right.example/v1".to_string()),
1089 route_decision: None,
1090 usage: Some(UsageMetrics {
1091 input_tokens: 10,
1092 output_tokens: 20,
1093 total_tokens: 30,
1094 ..UsageMetrics::default()
1095 }),
1096 cost: CostBreakdown::unknown(),
1097 retry: None,
1098 observability: Default::default(),
1099 service: "codex".to_string(),
1100 method: "POST".to_string(),
1101 path: "/v1/responses".to_string(),
1102 status_code: 429,
1103 duration_ms: 120,
1104 ttfb_ms: Some(20),
1105 streaming: false,
1106 ended_at_ms: 950,
1107 }];
1108 let mut balances = HashMap::new();
1109 balances.insert(
1110 "station-a".to_string(),
1111 vec![ProviderBalanceSnapshot {
1112 provider_id: "right".to_string(),
1113 station_name: Some("station-a".to_string()),
1114 upstream_index: Some(0),
1115 exhausted: Some(true),
1116 fetched_at_ms: 940,
1117 ..ProviderBalanceSnapshot::default()
1118 }],
1119 );
1120 let explain = RoutingExplainResponse {
1121 api_version: 1,
1122 service_name: "codex".to_string(),
1123 runtime_loaded_at_ms: None,
1124 request_model: None,
1125 session_id: None,
1126 request_context: Default::default(),
1127 selected_route: None,
1128 candidates: vec![RoutingExplainCandidate {
1129 provider_id: "right".to_string(),
1130 provider_alias: None,
1131 endpoint_id: "default".to_string(),
1132 provider_endpoint_key: "codex:right:default".to_string(),
1133 route_path: vec!["main".to_string(), "right".to_string()],
1134 preference_group: 0,
1135 compatibility: Some(crate::routing_explain::RoutingExplainCompatibility {
1136 station_name: "station-a".to_string(),
1137 upstream_index: 0,
1138 }),
1139 upstream_base_url: "https://right.example/v1".to_string(),
1140 selected: false,
1141 skip_reasons: vec![RoutingExplainSkipReason::UsageExhausted],
1142 }],
1143 affinity_policy: "off".to_string(),
1144 affinity: None,
1145 conditional_routes: Vec::new(),
1146 };
1147
1148 let view = build_view(
1149 &UsageRollupView::default(),
1150 &balances,
1151 &recent,
1152 Some(&explain),
1153 );
1154 let endpoint = view
1155 .endpoint_rows
1156 .iter()
1157 .find(|row| row.provider_id == "right")
1158 .expect("right endpoint");
1159
1160 assert_eq!(endpoint.endpoint_id, "default");
1161 assert_eq!(endpoint.usage.requests_total, 1);
1162 assert_eq!(endpoint.usage.requests_error, 1);
1163 assert_eq!(endpoint.balance_status, UsageBalanceStatus::Exhausted);
1164 assert_eq!(endpoint.route_skip_reasons, vec!["usage_exhausted"]);
1165 }
1166}