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