1use crate::config::v2::{Budget, BudgetTactic, BudgetWindow, WindowUnit};
26use crate::wire::intel::Usage;
27use serde::{Deserialize, Serialize};
28use serde_json::{Value, json};
29use std::collections::BTreeMap;
30
31#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum Admission {
34 Ok {
36 reservation: u64,
37 model: Option<String>,
38 },
39 Wait { until_ms: u64, reason: String },
41 Refuse { reason: String },
43 Fail { reason: String },
45}
46
47#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
49pub struct WindowState {
50 pub index: u64,
51 pub tokens: u64,
52 pub requests: u64,
53 #[serde(skip)]
54 pub reserved: u64,
55}
56
57#[derive(Debug, Clone)]
59struct Window {
60 cfg: BudgetWindow,
61 reset_offset_ms: u64,
63 state: WindowState,
64}
65
66impl Window {
67 fn len_ms(&self) -> u64 {
68 self.cfg.per.duration().as_millis() as u64
69 }
70 fn index_at(&self, now_ms: u64) -> u64 {
72 match self.cfg.per {
73 WindowUnit::Day | WindowUnit::Week => {
74 now_ms.saturating_sub(self.reset_offset_ms + week_epoch_shift(self.cfg.per))
75 / self.len_ms()
76 }
77 _ => now_ms / self.len_ms(),
78 }
79 }
80 fn start_ms(&self, now_ms: u64) -> u64 {
81 let idx = self.index_at(now_ms);
82 match self.cfg.per {
83 WindowUnit::Day | WindowUnit::Week => {
84 idx * self.len_ms() + self.reset_offset_ms + week_epoch_shift(self.cfg.per)
85 }
86 _ => idx * self.len_ms(),
87 }
88 }
89 fn next_reset_ms(&self, now_ms: u64) -> u64 {
90 self.start_ms(now_ms) + self.len_ms()
91 }
92 fn roll(&mut self, now_ms: u64) {
94 let idx = self.index_at(now_ms);
95 if idx != self.state.index {
96 self.state = WindowState {
97 index: idx,
98 ..Default::default()
99 };
100 }
101 }
102 fn tokens_left(&self) -> Option<u64> {
103 self.cfg
104 .tokens
105 .map(|cap| cap.saturating_sub(self.state.tokens + self.state.reserved))
106 }
107 fn requests_left(&self) -> Option<u64> {
108 self.cfg
109 .requests
110 .map(|cap| cap.saturating_sub(self.state.requests))
111 }
112 fn label(&self) -> String {
113 format!("{:?}", self.cfg.per).to_lowercase()
114 }
115}
116
117fn week_epoch_shift(unit: WindowUnit) -> u64 {
119 match unit {
120 WindowUnit::Week => 4 * 86_400_000,
123 _ => 0,
124 }
125}
126
127pub fn parse_reset(s: &str) -> Option<u64> {
129 let t = s.trim().trim_end_matches(['Z', 'z']);
130 let (h, m) = t.split_once(':')?;
131 let h: u64 = h.parse().ok()?;
132 let m: u64 = m.parse().ok()?;
133 (h < 24 && m < 60).then_some((h * 3600 + m * 60) * 1000)
134}
135
136#[derive(Debug, Clone)]
138struct Scope {
139 windows: Vec<Window>,
140 lifetime_cap: u64,
141 lifetime_used: u64,
142 tactic: BudgetTactic,
143 slow_factor: f64,
144 degrade_model: Option<String>,
145}
146
147impl Scope {
148 fn from_budget(b: &Budget) -> Scope {
149 Scope {
150 windows: b
151 .windows
152 .iter()
153 .map(|w| Window {
154 cfg: w.clone(),
155 reset_offset_ms: w.reset.as_deref().and_then(parse_reset).unwrap_or(0),
156 state: WindowState::default(),
157 })
158 .collect(),
159 lifetime_cap: b.lifetime_tokens.unwrap_or(0),
160 lifetime_used: 0,
161 tactic: b.on_exhausted,
162 slow_factor: b.slow.factor.unwrap_or(0.5).clamp(0.01, 1.0),
163 degrade_model: b.degrade.model.clone(),
164 }
165 }
166
167 fn roll(&mut self, now_ms: u64) {
168 for w in &mut self.windows {
169 w.roll(now_ms);
170 }
171 }
172
173 fn check(&self, estimate: u64, now_ms: u64) -> Result<(), Exhausted> {
175 if self.lifetime_cap > 0 && self.lifetime_used + estimate > self.lifetime_cap {
176 return Err(Exhausted {
177 window: "lifetime".into(),
178 until_ms: None,
179 pacing: false,
180 });
181 }
182 for w in &self.windows {
183 if let Some(left) = w.tokens_left()
184 && estimate > left
185 {
186 return Err(Exhausted {
187 window: w.label(),
188 until_ms: Some(w.next_reset_ms(now_ms)),
189 pacing: false,
190 });
191 }
192 if let Some(left) = w.requests_left()
193 && left == 0
194 {
195 return Err(Exhausted {
196 window: w.label(),
197 until_ms: Some(w.next_reset_ms(now_ms)),
198 pacing: false,
199 });
200 }
201 if self.tactic == BudgetTactic::Slow
203 && let Some(cap) = w.cfg.tokens
204 {
205 let elapsed = now_ms.saturating_sub(w.start_ms(now_ms)) as f64;
206 let len = w.len_ms() as f64;
207 let allowed = self.slow_factor * cap as f64 * (elapsed / len).clamp(0.0, 1.0)
208 + self.slow_factor * cap as f64 * 0.05;
209 let after = (w.state.tokens + w.state.reserved + estimate) as f64;
210 if after > allowed && after <= cap as f64 {
211 let t = (after / (self.slow_factor * cap as f64)) * len;
213 let until = w.start_ms(now_ms) + t.min(len) as u64;
214 return Err(Exhausted {
215 window: w.label(),
216 until_ms: Some(until.max(now_ms + 50)),
217 pacing: true,
218 });
219 }
220 }
221 }
222 Ok(())
223 }
224
225 fn reserve(&mut self, estimate: u64) {
226 for w in &mut self.windows {
227 w.state.reserved += estimate;
228 w.state.requests += 1;
229 }
230 }
231
232 fn settle(&mut self, reserved: u64, used: u64) {
233 for w in &mut self.windows {
234 w.state.reserved = w.state.reserved.saturating_sub(reserved);
235 w.state.tokens += used;
236 }
237 self.lifetime_used += used;
238 }
239
240 fn release(&mut self, reserved: u64) {
241 for w in &mut self.windows {
242 w.state.reserved = w.state.reserved.saturating_sub(reserved);
243 w.state.requests = w.state.requests.saturating_sub(1);
244 }
245 }
246
247 fn to_value(&self) -> Value {
248 json!({
249 "windows": self.windows.iter().map(|w| json!({"per": w.label(), "index": w.state.index, "tokens": w.state.tokens, "requests": w.state.requests})).collect::<Vec<_>>(),
250 "lifetime_used": self.lifetime_used,
251 })
252 }
253
254 fn adopt(&mut self, v: &Value) {
255 if let Some(ws) = v.get("windows").and_then(Value::as_array) {
256 for w in &mut self.windows {
257 if let Some(saved) = ws.iter().find(|s| s["per"].as_str() == Some(&w.label())) {
258 w.state = WindowState {
259 index: saved["index"].as_u64().unwrap_or(0),
260 tokens: saved["tokens"].as_u64().unwrap_or(0),
261 requests: saved["requests"].as_u64().unwrap_or(0),
262 reserved: 0,
263 };
264 }
265 }
266 }
267 self.lifetime_used = v.get("lifetime_used").and_then(Value::as_u64).unwrap_or(0);
268 }
269
270 fn status(&self, now_ms: u64) -> Value {
271 json!({
272 "tactic": format!("{:?}", self.tactic).to_lowercase(),
273 "lifetime": if self.lifetime_cap > 0 { json!({"cap": self.lifetime_cap, "used": self.lifetime_used, "remaining": self.lifetime_cap.saturating_sub(self.lifetime_used)}) } else { Value::Null },
274 "windows": self.windows.iter().map(|w| json!({
275 "per": w.label(), "tokens": {"cap": w.cfg.tokens, "used": w.state.tokens, "reserved": w.state.reserved, "remaining": w.tokens_left()},
276 "requests": {"cap": w.cfg.requests, "used": w.state.requests, "remaining": w.requests_left()},
277 "resets_at_ms": w.next_reset_ms(now_ms),
278 })).collect::<Vec<_>>(),
279 })
280 }
281}
282
283#[derive(Debug, Clone, PartialEq, Eq)]
284struct Exhausted {
285 window: String,
286 until_ms: Option<u64>,
287 pacing: bool,
288}
289
290#[derive(Debug, Clone)]
292struct Reservation {
293 estimate: u64,
294 scopes: Vec<String>,
295}
296
297#[derive(Debug, Clone)]
299pub struct Governor {
300 instance: Scope,
301 scopes: BTreeMap<String, Scope>,
302 reservations: BTreeMap<u64, Reservation>,
303 next_reservation: u64,
304 pub waiting: BTreeMap<String, u64>,
306 events: u64,
307}
308
309impl Governor {
310 pub fn new(budget: &Budget) -> Governor {
311 Governor {
312 instance: Scope::from_budget(budget),
313 scopes: BTreeMap::new(),
314 reservations: BTreeMap::new(),
315 next_reservation: 1,
316 waiting: BTreeMap::new(),
317 events: 0,
318 }
319 }
320
321 pub fn is_active(&self) -> bool {
323 !self.instance.windows.is_empty()
324 || self.instance.lifetime_cap > 0
325 || !self.scopes.is_empty()
326 }
327
328 pub fn ensure_scope(&mut self, key: &str, budget: &Budget) {
330 self.scopes
331 .entry(key.to_string())
332 .or_insert_with(|| Scope::from_budget(budget));
333 }
334 pub fn drop_scope(&mut self, key: &str) {
335 self.scopes.remove(key);
336 }
337
338 pub fn admit(&mut self, estimate: u64, scopes: &[String], now_ms: u64) -> Admission {
342 self.instance.roll(now_ms);
343 for k in scopes {
344 if let Some(s) = self.scopes.get_mut(k) {
345 s.roll(now_ms);
346 }
347 }
348 let mut verdict: Option<(Exhausted, String, BudgetTactic, Option<String>)> = None;
352 for k in scopes {
353 if let Some(s) = self.scopes.get(k)
354 && let Err(ex) = s.check(estimate, now_ms)
355 {
356 verdict = Some((ex, k.clone(), s.tactic, s.degrade_model.clone()));
357 break;
358 }
359 }
360 if verdict.is_none()
361 && let Err(ex) = self.instance.check(estimate, now_ms)
362 {
363 verdict = Some((
364 ex,
365 "instance".into(),
366 self.instance.tactic,
367 self.instance.degrade_model.clone(),
368 ));
369 }
370 if let Some((ex, key, tactic, degrade_model)) = verdict {
371 let reason = if ex.pacing {
372 format!(
373 "budget pacing ({key} {} window): slowing admissions",
374 ex.window
375 )
376 } else {
377 format!("budget exhausted ({key} {} window)", ex.window)
378 };
379 self.events += 1;
380 if ex.window == "lifetime" {
381 return Admission::Fail {
382 reason: format!("lifetime token budget exhausted ({key})"),
383 };
384 }
385 return match tactic {
386 BudgetTactic::Wait | BudgetTactic::Slow => Admission::Wait {
387 until_ms: ex.until_ms.unwrap_or(now_ms + 1000),
388 reason,
389 },
390 BudgetTactic::Degrade => match degrade_model {
391 Some(m) => {
392 let id = self.reserve_all(estimate, scopes);
393 Admission::Ok {
394 reservation: id,
395 model: Some(m),
396 }
397 }
398 None => Admission::Wait {
399 until_ms: ex.until_ms.unwrap_or(now_ms + 1000),
400 reason,
401 },
402 },
403 BudgetTactic::Refuse => Admission::Refuse {
404 reason: format!("refused: {reason}"),
405 },
406 BudgetTactic::Fail => Admission::Fail { reason },
407 };
408 }
409 let id = self.reserve_all(estimate, scopes);
410 Admission::Ok {
411 reservation: id,
412 model: None,
413 }
414 }
415
416 fn reserve_all(&mut self, estimate: u64, scopes: &[String]) -> u64 {
417 self.instance.reserve(estimate);
418 for k in scopes {
419 if let Some(s) = self.scopes.get_mut(k) {
420 s.reserve(estimate);
421 }
422 }
423 let id = self.next_reservation;
424 self.next_reservation += 1;
425 self.reservations.insert(
426 id,
427 Reservation {
428 estimate,
429 scopes: scopes.to_vec(),
430 },
431 );
432 id
433 }
434
435 pub fn settle(&mut self, reservation: u64, usage: Usage) {
437 let Some(r) = self.reservations.remove(&reservation) else {
438 return;
439 };
440 let used = usage.total();
441 self.instance.settle(r.estimate, used);
442 for k in &r.scopes {
443 if let Some(s) = self.scopes.get_mut(k) {
444 s.settle(r.estimate, used);
445 }
446 }
447 }
448
449 pub fn charge(&mut self, usage: Usage, scopes: &[String]) {
452 let used = usage.total();
453 self.instance.settle(0, used);
454 for k in scopes {
455 if let Some(s) = self.scopes.get_mut(k) {
456 s.settle(0, used);
457 }
458 }
459 }
460
461 pub fn release(&mut self, reservation: u64) {
463 let Some(r) = self.reservations.remove(&reservation) else {
464 return;
465 };
466 self.instance.release(r.estimate);
467 for k in &r.scopes {
468 if let Some(s) = self.scopes.get_mut(k) {
469 s.release(r.estimate);
470 }
471 }
472 }
473
474 pub fn to_value(&self) -> Value {
476 json!({
477 "instance": self.instance.to_value(),
478 "scopes": self.scopes.iter().map(|(k, s)| (k.clone(), s.to_value())).collect::<BTreeMap<_, _>>(),
479 })
480 }
481
482 pub fn restore(&mut self, v: &Value, now_ms: u64) {
485 if let Some(i) = v.get("instance") {
486 self.instance.adopt(i);
487 }
488 if let Some(sc) = v.get("scopes").and_then(Value::as_object) {
489 for (k, sv) in sc {
490 if let Some(s) = self.scopes.get_mut(k) {
491 s.adopt(sv);
492 }
493 }
494 }
495 self.instance.roll(now_ms);
496 }
497
498 pub fn status(&self, now_ms: u64) -> Value {
500 json!({
501 "active": self.is_active(),
502 "instance": self.instance.status(now_ms),
503 "scopes": self.scopes.iter().map(|(k, s)| (k.clone(), s.status(now_ms))).collect::<BTreeMap<_, _>>(),
504 "reservations": self.reservations.len(),
505 "waiting": self.waiting,
506 "events": self.events,
507 })
508 }
509
510 pub fn lifetime_used(&self) -> u64 {
512 self.instance.lifetime_used
513 }
514
515 pub fn next_reset_ms(&self, now_ms: u64) -> Option<u64> {
517 self.instance
518 .windows
519 .iter()
520 .map(|w| w.next_reset_ms(now_ms))
521 .min()
522 }
523}
524
525#[cfg(test)]
526mod tests {
527 use super::*;
528
529 fn budget(doc: Value) -> Budget {
530 serde_json::from_value(doc).unwrap()
531 }
532 fn usage(n: u64) -> Usage {
533 Usage {
534 input_tokens: n,
535 output_tokens: 0,
536 }
537 }
538
539 #[test]
540 fn windows_reserve_settle_and_roll_over() {
541 let mut g = Governor::new(&budget(
542 json!({"windows": [{"per": "minute", "tokens": 1000, "requests": 3}], "on_exhausted": "wait"}),
543 ));
544 assert!(g.is_active());
545 let t0 = 1_700_000_000_000u64; let t0 = t0 - t0 % 60_000; let a = g.admit(400, &[], t0);
548 let Admission::Ok {
549 reservation: r1,
550 model: None,
551 } = a
552 else {
553 panic!("{a:?}")
554 };
555 assert!(
557 matches!(g.admit(700, &[], t0 + 1000), Admission::Wait { until_ms, .. } if until_ms == t0 + 60_000)
558 );
559 g.settle(r1, usage(300)); let Admission::Ok {
561 reservation: r2, ..
562 } = g.admit(700, &[], t0 + 2000)
563 else {
564 panic!()
565 };
566 g.settle(r2, usage(700));
567 assert!(matches!(g.admit(1, &[], t0 + 3000), Admission::Wait { .. }));
569 let Admission::Ok {
571 reservation: r3, ..
572 } = g.admit(900, &[], t0 + 60_000)
573 else {
574 panic!()
575 };
576 g.release(r3);
577 for _ in 0..3 {
579 let Admission::Ok { reservation, .. } = g.admit(1, &[], t0 + 61_000) else {
580 panic!()
581 };
582 g.settle(reservation, usage(1));
583 }
584 assert!(
585 matches!(g.admit(1, &[], t0 + 62_000), Admission::Wait { .. }),
586 "requests exhausted"
587 );
588 let v = g.to_value();
590 assert_eq!(v["instance"]["windows"][0]["tokens"], json!(3));
591 let mut g2 = Governor::new(&budget(
592 json!({"windows": [{"per": "minute", "tokens": 1000, "requests": 3}]}),
593 ));
594 g2.restore(&v, t0 + 63_000);
595 assert!(
596 matches!(g2.admit(1, &[], t0 + 63_000), Admission::Wait { .. }),
597 "restored counters still exhausted"
598 );
599 g2.restore(&v, t0 + 120_000);
600 assert!(
601 matches!(g2.admit(1, &[], t0 + 120_000), Admission::Ok { .. }),
602 "rolled to a new window"
603 );
604 let st = g2.status(t0 + 120_000);
605 assert_eq!(st["instance"]["windows"][0]["per"], json!("minute"));
606 }
607
608 #[test]
609 fn tactics_lifetime_and_calendar_windows() {
610 let day = 86_400_000u64;
611 let now = 1_700_000_000_000u64;
612 let mut g = Governor::new(&budget(
614 json!({"windows": [{"per": "day", "tokens": 100, "reset": "06:00Z"}], "on_exhausted": "fail"}),
615 ));
616 let start_of_day = now - now % day;
617 let before = start_of_day + 5 * 3_600_000; let after = start_of_day + 7 * 3_600_000; let Admission::Ok { reservation, .. } = g.admit(100, &[], before) else {
620 panic!()
621 };
622 g.settle(reservation, usage(100));
623 assert!(
624 matches!(g.admit(1, &[], before + 60_000), Admission::Fail { .. }),
625 "fail tactic"
626 );
627 assert!(
628 matches!(g.admit(1, &[], after), Admission::Ok { .. }),
629 "the 06:00Z reset opened a new day window"
630 );
631 let mut g = Governor::new(&budget(
633 json!({"windows": [{"per": "hour", "tokens": 10}], "on_exhausted": "refuse"}),
634 ));
635 assert!(matches!(g.admit(11, &[], now), Admission::Refuse { .. }));
636 let mut g = Governor::new(&budget(
638 json!({"windows": [{"per": "hour", "tokens": 10}], "on_exhausted": "degrade", "degrade": {"model": "cheap"}}),
639 ));
640 assert!(
641 matches!(g.admit(11, &[], now), Admission::Ok { model: Some(m), .. } if m == "cheap")
642 );
643 let mut g = Governor::new(&budget(
645 json!({"windows": [{"per": "hour", "tokens": 3600}], "on_exhausted": "slow", "slow": {"factor": 0.5}}),
646 ));
647 let hour_start = now - now % 3_600_000;
648 assert!(matches!(g.admit(80, &[], hour_start), Admission::Ok { .. }));
650 assert!(
651 matches!(g.admit(500, &[], hour_start + 1000), Admission::Wait { until_ms, .. } if until_ms > hour_start + 1000 && until_ms < hour_start + 3_600_000)
652 );
653 assert!(matches!(
655 g.admit(500, &[], hour_start + 1_800_000),
656 Admission::Ok { .. }
657 ));
658 let mut g = Governor::new(&budget(
660 json!({"lifetime_tokens": 50, "on_exhausted": "wait"}),
661 ));
662 let Admission::Ok { reservation, .. } = g.admit(30, &[], now) else {
663 panic!()
664 };
665 g.settle(reservation, usage(30));
666 assert!(
667 matches!(g.admit(30, &[], now), Admission::Fail { reason } if reason.contains("lifetime"))
668 );
669 assert_eq!(g.lifetime_used(), 30);
670 let mut g = Governor::new(&budget(json!({})));
672 assert!(!g.is_active());
673 assert!(matches!(g.admit(1_000_000, &[], now), Admission::Ok { .. }));
674 }
675
676 #[test]
677 fn sub_scopes_nest_under_the_instance() {
678 let mut g = Governor::new(&budget(
679 json!({"windows": [{"per": "hour", "tokens": 1000}], "on_exhausted": "wait"}),
680 ));
681 g.ensure_scope(
682 "conversation:c1",
683 &budget(json!({"windows": [{"per": "hour", "tokens": 100}], "on_exhausted": "refuse"})),
684 );
685 let now = 1_700_000_000_000u64;
686 assert!(
688 matches!(g.admit(150, &["conversation:c1".to_string()], now), Admission::Refuse { reason } if reason.contains("conversation:c1"))
689 );
690 let Admission::Ok { reservation, .. } = g.admit(50, &["conversation:c1".to_string()], now)
692 else {
693 panic!()
694 };
695 g.settle(reservation, usage(50));
696 let v = g.to_value();
697 assert_eq!(
698 v["scopes"]["conversation:c1"]["windows"][0]["tokens"],
699 json!(50)
700 );
701 assert_eq!(v["instance"]["windows"][0]["tokens"], json!(50));
702 g.charge(usage(940), &[]);
704 assert!(matches!(g.admit(20, &[], now), Admission::Wait { .. }));
705 assert_eq!(parse_reset("06:30Z"), Some((6 * 3600 + 30 * 60) * 1000));
706 assert_eq!(parse_reset("25:00Z"), None);
707 g.drop_scope("conversation:c1");
708 assert!(g.to_value()["scopes"].as_object().unwrap().is_empty());
709 }
710}