1use serde::{Deserialize, Serialize};
2
3use crate::balance::ProviderBalanceSnapshot;
4use crate::config::UsageForecastConfig;
5use crate::pricing::UsdAmount;
6use crate::state::FinishedRequest;
7
8const MINUTE_MS: u64 = 60_000;
9const HOUR_MS: u64 = 60 * MINUTE_MS;
10const DAY_MS: u64 = 24 * HOUR_MS;
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
13pub struct UsageSpendForecast {
14 pub enabled: bool,
15 pub sample_window_ms: u64,
16 pub sample_elapsed_ms: u64,
17 pub projected_sample_elapsed_ms: u64,
18 pub priced_requests: u64,
19 pub unpriced_requests: u64,
20 pub sample_cost_usd: Option<String>,
21 pub rate_per_hour_usd: Option<String>,
22 pub reset_at_ms: Option<u64>,
23 pub reset_in_ms: Option<u64>,
24 pub projected_until_reset_usd: Option<String>,
25 pub primary_balance_remaining_usd: Option<String>,
26 pub projected_balance_after_reset_usd: Option<String>,
27 pub projected_exhaustion: bool,
28 pub confidence: UsageForecastConfidence,
29 pub reason: Option<String>,
30}
31
32#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
33#[serde(rename_all = "snake_case")]
34pub enum UsageForecastConfidence {
35 #[default]
36 Disabled,
37 NoData,
38 LowSample,
39 PartialPricing,
40 Estimated,
41}
42
43pub fn build_usage_spend_forecast(
44 config: &UsageForecastConfig,
45 recent: &[FinishedRequest],
46 provider_balances: &[ProviderBalanceSnapshot],
47 now_ms: u64,
48) -> UsageSpendForecast {
49 if !config.enabled {
50 return UsageSpendForecast {
51 enabled: false,
52 reason: Some("disabled".to_string()),
53 ..UsageSpendForecast::default()
54 };
55 }
56
57 let sample_window_ms = config.rate_window_minutes.max(1).saturating_mul(MINUTE_MS);
58 let reset_at_ms = next_reset_at_ms(
59 now_ms,
60 config.reset_utc_offset.as_str(),
61 config.reset_time.as_str(),
62 );
63 let reset_in_ms = reset_at_ms.map(|reset| reset.saturating_sub(now_ms));
64
65 let cutoff_ms = now_ms.saturating_sub(sample_window_ms);
66 let mut priced_requests = 0_u64;
67 let mut unpriced_requests = 0_u64;
68 let mut sample_cost = UsdAmount::ZERO;
69 let mut oldest_sample_ms: Option<u64> = None;
70
71 for request in recent {
72 if request.ended_at_ms < cutoff_ms || request.ended_at_ms > now_ms {
73 continue;
74 }
75 if request.cost.total_cost_usd.is_some() {
76 if let Some(cost) = request
77 .cost
78 .total_cost_usd
79 .as_deref()
80 .and_then(UsdAmount::from_decimal_str)
81 {
82 sample_cost = sample_cost.saturating_add(cost);
83 priced_requests = priced_requests.saturating_add(1);
84 oldest_sample_ms = Some(
85 oldest_sample_ms
86 .map(|oldest| oldest.min(request.ended_at_ms))
87 .unwrap_or(request.ended_at_ms),
88 );
89 } else {
90 unpriced_requests = unpriced_requests.saturating_add(1);
91 }
92 } else if request.usage.is_some() {
93 unpriced_requests = unpriced_requests.saturating_add(1);
94 }
95 }
96
97 let sample_elapsed_ms = oldest_sample_ms
98 .map(|oldest| {
99 now_ms
100 .saturating_sub(oldest)
101 .clamp(MINUTE_MS, sample_window_ms)
102 })
103 .unwrap_or(0);
104
105 if priced_requests == 0 || sample_elapsed_ms == 0 {
106 return UsageSpendForecast {
107 enabled: true,
108 sample_window_ms,
109 sample_elapsed_ms,
110 priced_requests,
111 unpriced_requests,
112 reset_at_ms,
113 reset_in_ms,
114 confidence: UsageForecastConfidence::NoData,
115 reason: Some("no priced requests in forecast window".to_string()),
116 ..UsageSpendForecast::default()
117 };
118 }
119
120 let rate_per_hour = scale_usd_by_ratio(sample_cost, HOUR_MS, sample_elapsed_ms);
121 let confidence = if priced_requests < config.min_priced_requests.max(1) {
122 UsageForecastConfidence::LowSample
123 } else if unpriced_requests > 0 {
124 UsageForecastConfidence::PartialPricing
125 } else {
126 UsageForecastConfidence::Estimated
127 };
128 let projected_sample_elapsed_ms = if confidence == UsageForecastConfidence::LowSample {
129 0
130 } else {
131 sample_elapsed_ms
132 };
133 let projected_until_reset = if confidence == UsageForecastConfidence::LowSample {
134 None
135 } else {
136 reset_in_ms.map(|reset_in_ms| scale_usd_by_ratio(rate_per_hour, reset_in_ms, HOUR_MS))
137 };
138 let primary_balance_remaining =
139 primary_remaining_balance(provider_balances, now_ms).map(|(_, amount)| amount);
140 let projected_balance_after_reset = match (primary_balance_remaining, projected_until_reset) {
141 (Some(balance), Some(projected)) => Some(balance.saturating_sub(projected)),
142 _ => None,
143 };
144 let projected_exhaustion = matches!(
145 (primary_balance_remaining, projected_until_reset),
146 (Some(balance), Some(projected)) if projected > balance
147 );
148
149 UsageSpendForecast {
150 enabled: true,
151 sample_window_ms,
152 sample_elapsed_ms,
153 projected_sample_elapsed_ms,
154 priced_requests,
155 unpriced_requests,
156 sample_cost_usd: Some(sample_cost.format_usd()),
157 rate_per_hour_usd: Some(rate_per_hour.format_usd()),
158 reset_at_ms,
159 reset_in_ms,
160 projected_until_reset_usd: projected_until_reset.map(UsdAmount::format_usd),
161 primary_balance_remaining_usd: primary_balance_remaining.map(UsdAmount::format_usd),
162 projected_balance_after_reset_usd: projected_balance_after_reset.map(UsdAmount::format_usd),
163 projected_exhaustion,
164 confidence,
165 reason: None,
166 }
167}
168
169fn scale_usd_by_ratio(amount: UsdAmount, numerator: u64, denominator: u64) -> UsdAmount {
170 if denominator == 0 || numerator == 0 || amount.is_zero() {
171 return UsdAmount::ZERO;
172 }
173 UsdAmount::from_femto_usd(
174 amount
175 .femto_usd()
176 .saturating_mul(i128::from(numerator))
177 .saturating_div(i128::from(denominator)),
178 )
179}
180
181fn primary_remaining_balance(
182 provider_balances: &[ProviderBalanceSnapshot],
183 now_ms: u64,
184) -> Option<(&ProviderBalanceSnapshot, UsdAmount)> {
185 provider_balances
186 .iter()
187 .filter(|snapshot| snapshot.error.as_deref().is_none_or(str::is_empty))
188 .filter_map(|snapshot| remaining_balance(snapshot).map(|amount| (snapshot, amount)))
189 .min_by(|(left, _), (right, _)| {
190 balance_rank(left, now_ms)
191 .cmp(&balance_rank(right, now_ms))
192 .then_with(|| left.upstream_index.cmp(&right.upstream_index))
193 .then_with(|| right.fetched_at_ms.cmp(&left.fetched_at_ms))
194 })
195}
196
197fn remaining_balance(snapshot: &ProviderBalanceSnapshot) -> Option<UsdAmount> {
198 snapshot
199 .quota_remaining_usd
200 .as_deref()
201 .and_then(UsdAmount::from_decimal_str)
202 .or_else(|| {
203 snapshot
204 .total_balance_usd
205 .as_deref()
206 .and_then(UsdAmount::from_decimal_str)
207 })
208}
209
210fn balance_rank(snapshot: &ProviderBalanceSnapshot, now_ms: u64) -> u8 {
211 match snapshot.status_at(now_ms) {
212 crate::balance::BalanceSnapshotStatus::Ok => 0,
213 crate::balance::BalanceSnapshotStatus::Stale => 1,
214 crate::balance::BalanceSnapshotStatus::Unknown => 2,
215 crate::balance::BalanceSnapshotStatus::Error => 3,
216 crate::balance::BalanceSnapshotStatus::Exhausted => 4,
217 }
218}
219
220pub fn next_reset_at_ms(now_ms: u64, utc_offset: &str, reset_time: &str) -> Option<u64> {
221 let offset_ms = parse_utc_offset_ms(utc_offset)?;
222 let reset_ms = parse_hh_mm_ms(reset_time)?;
223 let local_ms = i128::from(now_ms) + offset_ms;
224 let local_day_start = div_floor(local_ms, i128::from(DAY_MS)) * i128::from(DAY_MS);
225 let mut reset_local = local_day_start + i128::from(reset_ms);
226 if reset_local <= local_ms {
227 reset_local += i128::from(DAY_MS);
228 }
229 let reset_utc = reset_local - offset_ms;
230 u64::try_from(reset_utc).ok()
231}
232
233fn parse_hh_mm_ms(value: &str) -> Option<u64> {
234 let (hour, minute) = value.trim().split_once(':')?;
235 let hour = hour.parse::<u64>().ok()?;
236 let minute = minute.parse::<u64>().ok()?;
237 if hour >= 24 || minute >= 60 {
238 return None;
239 }
240 Some(
241 hour.saturating_mul(HOUR_MS)
242 .saturating_add(minute.saturating_mul(MINUTE_MS)),
243 )
244}
245
246fn parse_utc_offset_ms(value: &str) -> Option<i128> {
247 let value = value.trim();
248 if value.eq_ignore_ascii_case("z") || value == "+00:00" || value == "-00:00" {
249 return Some(0);
250 }
251 let sign = match value.as_bytes().first().copied()? {
252 b'+' => 1_i128,
253 b'-' => -1_i128,
254 _ => return None,
255 };
256 let rest = &value[1..];
257 let (hour, minute) = rest.split_once(':')?;
258 let hour = hour.parse::<i128>().ok()?;
259 let minute = minute.parse::<i128>().ok()?;
260 if hour > 23 || minute > 59 {
261 return None;
262 }
263 Some(sign * (hour * i128::from(HOUR_MS) + minute * i128::from(MINUTE_MS)))
264}
265
266fn div_floor(a: i128, b: i128) -> i128 {
267 let q = a / b;
268 let r = a % b;
269 if r != 0 && ((r > 0) != (b > 0)) {
270 q - 1
271 } else {
272 q
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use super::*;
279 use crate::pricing::{CostAdjustments, ModelPrice, estimate_usage_cost_with_accounting};
280 use crate::state::RequestObservability;
281 use crate::usage::{CacheInputAccounting, UsageMetrics};
282
283 fn priced_request(ended_at_ms: u64, usd: &str) -> FinishedRequest {
284 let usage = UsageMetrics {
285 input_tokens: 1_000_000,
286 total_tokens: 1_000_000,
287 ..Default::default()
288 };
289 let price = ModelPrice::from_per_million_usd(
290 "gpt-test",
291 None,
292 usd,
293 "0",
294 Some("0"),
295 Some("0"),
296 "test",
297 )
298 .expect("test price");
299 let cost = estimate_usage_cost_with_accounting(
300 &usage,
301 &price,
302 CostAdjustments::default(),
303 CacheInputAccounting::default(),
304 );
305
306 FinishedRequest {
307 id: ended_at_ms,
308 trace_id: None,
309 session_id: None,
310 session_identity_source: None,
311 client_name: None,
312 client_addr: None,
313 cwd: None,
314 model: Some("gpt-test".to_string()),
315 reasoning_effort: None,
316 service_tier: None,
317 station_name: Some("station".to_string()),
318 provider_id: Some("provider".to_string()),
319 upstream_base_url: None,
320 route_decision: None,
321 usage: Some(usage),
322 cost,
323 retry: None,
324 observability: RequestObservability::default(),
325 service: "codex".to_string(),
326 method: "POST".to_string(),
327 path: "/v1/responses".to_string(),
328 status_code: 200,
329 duration_ms: 100,
330 ttfb_ms: None,
331 streaming: false,
332 ended_at_ms,
333 }
334 }
335
336 #[test]
337 fn next_reset_uses_configured_fixed_offset_midnight() {
338 let now_ms = 10 * HOUR_MS;
339
340 let reset = next_reset_at_ms(now_ms, "+08:00", "00:00").expect("reset");
341
342 assert_eq!(reset, 16 * HOUR_MS);
343 }
344
345 #[test]
346 fn forecast_projects_current_hourly_spend_until_reset() {
347 let now_ms = 10 * HOUR_MS;
348 let config = UsageForecastConfig {
349 rate_window_minutes: 60,
350 min_priced_requests: 2,
351 ..UsageForecastConfig::default()
352 };
353 let recent = vec![
354 priced_request(now_ms - 60 * MINUTE_MS, "1"),
355 priced_request(now_ms - 30 * MINUTE_MS, "1"),
356 ];
357 let balances = vec![ProviderBalanceSnapshot {
358 provider_id: "provider".to_string(),
359 quota_remaining_usd: Some("20".to_string()),
360 fetched_at_ms: now_ms,
361 ..ProviderBalanceSnapshot::default()
362 }];
363
364 let forecast = build_usage_spend_forecast(&config, &recent, &balances, now_ms);
365
366 assert_eq!(forecast.rate_per_hour_usd.as_deref(), Some("2"));
367 assert_eq!(forecast.reset_in_ms, Some(6 * HOUR_MS));
368 assert_eq!(forecast.projected_until_reset_usd.as_deref(), Some("12"));
369 assert_eq!(
370 forecast.projected_balance_after_reset_usd.as_deref(),
371 Some("8")
372 );
373 assert!(!forecast.projected_exhaustion);
374 assert_eq!(forecast.confidence, UsageForecastConfidence::Estimated);
375 }
376
377 #[test]
378 fn forecast_marks_exhaustion_when_projection_exceeds_remaining_balance() {
379 let now_ms = 10 * HOUR_MS;
380 let config = UsageForecastConfig {
381 min_priced_requests: 2,
382 ..UsageForecastConfig::default()
383 };
384 let recent = vec![
385 priced_request(now_ms - 60 * MINUTE_MS, "1"),
386 priced_request(now_ms - 30 * MINUTE_MS, "1"),
387 ];
388 let balances = vec![ProviderBalanceSnapshot {
389 provider_id: "provider".to_string(),
390 quota_remaining_usd: Some("2".to_string()),
391 fetched_at_ms: now_ms,
392 ..ProviderBalanceSnapshot::default()
393 }];
394
395 let forecast = build_usage_spend_forecast(&config, &recent, &balances, now_ms);
396
397 assert_eq!(forecast.projected_until_reset_usd.as_deref(), Some("12"));
398 assert_eq!(
399 forecast.projected_balance_after_reset_usd.as_deref(),
400 Some("0")
401 );
402 assert!(forecast.projected_exhaustion);
403 }
404
405 #[test]
406 fn forecast_keeps_burn_rate_but_suppresses_projection_for_low_sample() {
407 let now_ms = 10 * HOUR_MS;
408 let config = UsageForecastConfig {
409 min_priced_requests: 2,
410 ..UsageForecastConfig::default()
411 };
412 let recent = vec![priced_request(now_ms - 30 * MINUTE_MS, "1")];
413
414 let forecast = build_usage_spend_forecast(&config, &recent, &[], now_ms);
415
416 assert_eq!(forecast.rate_per_hour_usd.as_deref(), Some("2"));
417 assert_eq!(forecast.confidence, UsageForecastConfidence::LowSample);
418 assert_eq!(forecast.projected_until_reset_usd, None);
419 assert_eq!(forecast.projected_sample_elapsed_ms, 0);
420 assert!(!forecast.projected_exhaustion);
421 }
422}