1use std::collections::BTreeMap;
2use std::sync::Arc;
3use std::time::Instant;
4
5use serde::{Deserialize, Serialize};
6
7mod wire;
8
9use crate::types::{CacheUsageStatus, TokenUsage, UsageSource};
10
11pub const MAX_WIRE_INTEGER: u64 = (1_u64 << 53) - 1;
12
13#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum UnavailableMetricPolicy {
16 #[default]
17 ContinueAndMark,
18 Stop,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
22#[serde(rename_all = "snake_case")]
23pub enum BudgetDimension {
24 TotalTokens,
25 UncachedInputTokens,
26 ToolCalls,
27 ToolCallsByName,
28 WallTime,
29 HostCost,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33#[serde(rename_all = "snake_case")]
34pub enum BudgetEnforcementBoundary {
35 RunStart,
36 CycleStart,
37 LlmComplete,
38 ToolBatchPreflight,
39 ToolBatchComplete,
40 Terminal,
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "snake_case")]
45pub enum BudgetExhaustionReason {
46 LimitReached,
47 LimitExceeded,
48 MetricUnavailable,
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(rename_all = "snake_case")]
53pub enum BudgetUnavailableReason {
54 UsageMissing,
55 MeterMissing,
56 MeterUnavailable,
57 MeterError,
58 UnitMismatch,
59 CurrencyMismatch,
60 NonMonotonic,
61 IntegerOverflow,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
65pub struct HostCost {
66 pub unit: String,
67 pub currency: Option<String>,
68 pub amount_microunits: u64,
69}
70
71impl HostCost {
72 pub fn new(unit: impl Into<String>, amount_microunits: u64) -> Result<Self, String> {
73 let cost = Self {
74 unit: unit.into(),
75 currency: None,
76 amount_microunits,
77 };
78 cost.validate()?;
79 Ok(cost)
80 }
81
82 pub fn with_currency(mut self, currency: impl Into<String>) -> Result<Self, String> {
83 self.currency = Some(currency.into());
84 self.validate()?;
85 Ok(self)
86 }
87
88 fn validate(&self) -> Result<(), String> {
89 validate_non_empty(&self.unit, "host cost unit")?;
90 if let Some(currency) = &self.currency {
91 validate_non_empty(currency, "host cost currency")?;
92 }
93 validate_wire_integer(self.amount_microunits, "host cost amount_microunits")
94 }
95}
96
97pub trait HostCostMeter: Send + Sync {
98 fn read(&self) -> Result<Option<HostCost>, String>;
99}
100
101#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
102pub struct BudgetUnavailableDimension {
103 pub dimension: BudgetDimension,
104 pub reason: BudgetUnavailableReason,
105 pub expected_unit: Option<String>,
106 pub observed_unit: Option<String>,
107 pub expected_currency: Option<String>,
108 pub observed_currency: Option<String>,
109}
110
111impl BudgetUnavailableDimension {
112 fn validate(&self) -> Result<(), String> {
113 for (field, value) in [
114 ("expected_unit", self.expected_unit.as_deref()),
115 ("observed_unit", self.observed_unit.as_deref()),
116 ("expected_currency", self.expected_currency.as_deref()),
117 ("observed_currency", self.observed_currency.as_deref()),
118 ] {
119 if let Some(value) = value {
120 validate_non_empty(value, field)?;
121 }
122 }
123 Ok(())
124 }
125}
126
127#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
128pub struct RunBudgetLimits {
129 pub max_total_tokens: Option<u64>,
130 pub max_uncached_input_tokens: Option<u64>,
131 pub max_tool_calls: Option<u64>,
132 pub max_tool_calls_by_name: BTreeMap<String, u64>,
133 pub max_wall_time_ms: Option<u64>,
134 pub max_host_cost: Option<HostCost>,
135 pub unavailable_metric_policy: UnavailableMetricPolicy,
136}
137
138impl RunBudgetLimits {
139 pub fn builder() -> RunBudgetLimitsBuilder {
140 RunBudgetLimitsBuilder::default()
141 }
142
143 pub fn has_limits(&self) -> bool {
144 self.max_total_tokens.is_some()
145 || self.max_uncached_input_tokens.is_some()
146 || self.max_tool_calls.is_some()
147 || !self.max_tool_calls_by_name.is_empty()
148 || self.max_wall_time_ms.is_some()
149 || self.max_host_cost.is_some()
150 }
151
152 pub fn validate(&self) -> Result<(), String> {
153 for (field, value) in [
154 ("max_total_tokens", self.max_total_tokens),
155 ("max_uncached_input_tokens", self.max_uncached_input_tokens),
156 ("max_tool_calls", self.max_tool_calls),
157 ("max_wall_time_ms", self.max_wall_time_ms),
158 ] {
159 if let Some(value) = value {
160 validate_wire_integer(value, field)?;
161 }
162 }
163 for (name, value) in &self.max_tool_calls_by_name {
164 validate_non_empty(name, "named tool budget key")?;
165 validate_wire_integer(*value, &format!("max_tool_calls_by_name[{name:?}]"))?;
166 }
167 if let Some(cost) = &self.max_host_cost {
168 cost.validate()?;
169 }
170 Ok(())
171 }
172}
173
174#[derive(Debug, Default)]
175pub struct RunBudgetLimitsBuilder {
176 limits: RunBudgetLimits,
177}
178
179impl RunBudgetLimitsBuilder {
180 pub fn max_total_tokens(mut self, value: u64) -> Self {
181 self.limits.max_total_tokens = Some(value);
182 self
183 }
184
185 pub fn max_uncached_input_tokens(mut self, value: u64) -> Self {
186 self.limits.max_uncached_input_tokens = Some(value);
187 self
188 }
189
190 pub fn max_tool_calls(mut self, value: u64) -> Self {
191 self.limits.max_tool_calls = Some(value);
192 self
193 }
194
195 pub fn max_tool_calls_by_name(
196 mut self,
197 values: impl IntoIterator<Item = (impl Into<String>, u64)>,
198 ) -> Self {
199 self.limits.max_tool_calls_by_name = values
200 .into_iter()
201 .map(|(name, value)| (name.into(), value))
202 .collect();
203 self
204 }
205
206 pub fn max_wall_time_ms(mut self, value: u64) -> Self {
207 self.limits.max_wall_time_ms = Some(value);
208 self
209 }
210
211 pub fn max_host_cost(mut self, value: HostCost) -> Self {
212 self.limits.max_host_cost = Some(value);
213 self
214 }
215
216 pub fn unavailable_metric_policy(mut self, value: UnavailableMetricPolicy) -> Self {
217 self.limits.unavailable_metric_policy = value;
218 self
219 }
220
221 pub fn build(self) -> Result<RunBudgetLimits, String> {
222 self.limits.validate()?;
223 Ok(self.limits)
224 }
225}
226
227#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
228pub struct BudgetUsageSnapshot {
229 pub cycles: u64,
230 pub total_tokens: Option<u64>,
231 pub uncached_input_tokens: Option<u64>,
232 pub tool_calls: u64,
233 pub tool_calls_by_name: BTreeMap<String, u64>,
234 pub elapsed_ms: u64,
235 pub host_cost: Option<HostCost>,
236 pub unavailable_dimensions: Vec<BudgetUnavailableDimension>,
237}
238
239impl Default for BudgetUsageSnapshot {
240 fn default() -> Self {
241 Self {
242 cycles: 0,
243 total_tokens: Some(0),
244 uncached_input_tokens: Some(0),
245 tool_calls: 0,
246 tool_calls_by_name: BTreeMap::new(),
247 elapsed_ms: 0,
248 host_cost: None,
249 unavailable_dimensions: Vec::new(),
250 }
251 }
252}
253
254impl BudgetUsageSnapshot {
255 pub fn validate(&self) -> Result<(), String> {
256 validate_wire_integer(self.cycles, "budget usage cycles")?;
257 if let Some(value) = self.total_tokens {
258 validate_wire_integer(value, "budget usage total_tokens")?;
259 }
260 if let Some(value) = self.uncached_input_tokens {
261 validate_wire_integer(value, "budget usage uncached_input_tokens")?;
262 }
263 validate_wire_integer(self.tool_calls, "budget usage tool_calls")?;
264 validate_wire_integer(self.elapsed_ms, "budget usage elapsed_ms")?;
265 for (name, value) in &self.tool_calls_by_name {
266 validate_non_empty(name, "budget usage tool name")?;
267 validate_wire_integer(
268 *value,
269 &format!("budget usage tool_calls_by_name[{name:?}]"),
270 )?;
271 }
272 if let Some(cost) = &self.host_cost {
273 cost.validate()?;
274 }
275 let mut dimensions = self
276 .unavailable_dimensions
277 .iter()
278 .map(|observation| observation.dimension)
279 .collect::<Vec<_>>();
280 dimensions.sort();
281 dimensions.dedup();
282 if dimensions.len() != self.unavailable_dimensions.len() {
283 return Err("budget unavailable dimensions must be unique".to_string());
284 }
285 for observation in &self.unavailable_dimensions {
286 observation.validate()?;
287 }
288 Ok(())
289 }
290}
291
292#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
293pub struct BudgetExhaustion {
294 pub dimension: BudgetDimension,
295 pub tool_name: Option<String>,
296 pub reason: BudgetExhaustionReason,
297 pub limit: u64,
298 pub observed: Option<u64>,
299 pub attempted_increment: Option<u64>,
300 pub overshoot: Option<u64>,
301 pub unit: String,
302 pub currency: Option<String>,
303 pub enforcement_boundary: BudgetEnforcementBoundary,
304 pub unavailable_reason: Option<BudgetUnavailableReason>,
305}
306
307impl BudgetExhaustion {
308 fn validate(&self) -> Result<(), String> {
309 validate_wire_integer(self.limit, "budget exhaustion limit")?;
310 for (field, value) in [
311 ("observed", self.observed),
312 ("attempted_increment", self.attempted_increment),
313 ("overshoot", self.overshoot),
314 ] {
315 if let Some(value) = value {
316 validate_wire_integer(value, &format!("budget exhaustion {field}"))?;
317 }
318 }
319 validate_non_empty(&self.unit, "budget exhaustion unit")?;
320 if let Some(tool_name) = &self.tool_name {
321 validate_non_empty(tool_name, "tool_name")?;
322 }
323 if let Some(currency) = &self.currency {
324 validate_non_empty(currency, "currency")?;
325 }
326 Ok(())
327 }
328}
329
330pub type MonotonicClock = Arc<dyn Fn() -> u128 + Send + Sync>;
331
332pub struct BudgetEvaluator {
333 pub limits: RunBudgetLimits,
334 host_cost_meter: Option<Arc<dyn HostCostMeter>>,
335 clock_ns: MonotonicClock,
336 started_ns: u128,
337 cycles: u64,
338 total_tokens: Option<u64>,
339 uncached_input_tokens: Option<u64>,
340 tool_calls: u64,
341 tool_calls_by_name: BTreeMap<String, u64>,
342 base_elapsed_ms: u64,
343 elapsed_ms: u64,
344 host_cost: Option<HostCost>,
345 unavailable: BTreeMap<BudgetDimension, BudgetUnavailableDimension>,
346}
347
348impl BudgetEvaluator {
349 pub fn new(
350 limits: RunBudgetLimits,
351 host_cost_meter: Option<Arc<dyn HostCostMeter>>,
352 initial_usage: Option<BudgetUsageSnapshot>,
353 ) -> Result<Self, String> {
354 let started = Instant::now();
355 Self::with_clock(
356 limits,
357 host_cost_meter,
358 initial_usage,
359 Arc::new(move || started.elapsed().as_nanos()),
360 )
361 }
362
363 pub fn with_clock(
364 limits: RunBudgetLimits,
365 host_cost_meter: Option<Arc<dyn HostCostMeter>>,
366 initial_usage: Option<BudgetUsageSnapshot>,
367 clock_ns: MonotonicClock,
368 ) -> Result<Self, String> {
369 limits.validate()?;
370 if !limits.has_limits() {
371 return Err("BudgetEvaluator requires at least one configured limit".to_string());
372 }
373 let usage = initial_usage.unwrap_or_default();
374 usage.validate()?;
375 let unavailable = usage
376 .unavailable_dimensions
377 .iter()
378 .cloned()
379 .map(|item| (item.dimension, item))
380 .collect();
381 let started_ns = clock_ns();
382 Ok(Self {
383 limits,
384 host_cost_meter,
385 clock_ns,
386 started_ns,
387 cycles: usage.cycles,
388 total_tokens: usage.total_tokens,
389 uncached_input_tokens: usage.uncached_input_tokens,
390 tool_calls: usage.tool_calls,
391 tool_calls_by_name: usage.tool_calls_by_name,
392 base_elapsed_ms: usage.elapsed_ms,
393 elapsed_ms: usage.elapsed_ms,
394 host_cost: usage.host_cost,
395 unavailable,
396 })
397 }
398
399 pub fn snapshot(&self) -> BudgetUsageSnapshot {
400 let mut unavailable_dimensions = self.unavailable.values().cloned().collect::<Vec<_>>();
401 unavailable_dimensions.sort_by_key(|item| dimension_precedence(item.dimension));
402 BudgetUsageSnapshot {
403 cycles: self.cycles,
404 total_tokens: self.total_tokens,
405 uncached_input_tokens: self.uncached_input_tokens,
406 tool_calls: self.tool_calls,
407 tool_calls_by_name: self.tool_calls_by_name.clone(),
408 elapsed_ms: self.elapsed_ms,
409 host_cost: self.host_cost.clone(),
410 unavailable_dimensions,
411 }
412 }
413
414 pub fn run_start(&mut self) -> Option<BudgetExhaustion> {
415 let boundary = BudgetEnforcementBoundary::RunStart;
416 self.observe_boundary();
417 self.strict_unavailable(boundary).or_else(|| {
418 self.check_admission_limits(
419 boundary,
420 &[BudgetDimension::WallTime, BudgetDimension::HostCost],
421 )
422 })
423 }
424
425 pub fn cycle_start(&mut self) -> Option<BudgetExhaustion> {
426 let boundary = BudgetEnforcementBoundary::CycleStart;
427 self.observe_boundary();
428 self.strict_unavailable(boundary).or_else(|| {
429 self.check_admission_limits(
430 boundary,
431 &[
432 BudgetDimension::WallTime,
433 BudgetDimension::TotalTokens,
434 BudgetDimension::UncachedInputTokens,
435 BudgetDimension::HostCost,
436 ],
437 )
438 })
439 }
440
441 pub fn llm_complete(&mut self, usage: &TokenUsage) -> Option<BudgetExhaustion> {
442 let boundary = BudgetEnforcementBoundary::LlmComplete;
443 self.cycles = self.cycles.saturating_add(1);
444 self.observe_token_usage(usage);
445 self.observe_boundary();
446 self.strict_unavailable(boundary).or_else(|| {
447 self.check_exceeded_limits(
448 boundary,
449 &[
450 BudgetDimension::WallTime,
451 BudgetDimension::TotalTokens,
452 BudgetDimension::UncachedInputTokens,
453 BudgetDimension::HostCost,
454 ],
455 )
456 })
457 }
458
459 pub fn preflight_tools(&mut self, tool_names: &[String]) -> Option<BudgetExhaustion> {
460 let boundary = BudgetEnforcementBoundary::ToolBatchPreflight;
461 self.observe_boundary();
462 if let Some(exhaustion) = self.strict_unavailable(boundary).or_else(|| {
463 self.check_admission_limits(
464 boundary,
465 &[BudgetDimension::WallTime, BudgetDimension::HostCost],
466 )
467 }) {
468 return Some(exhaustion);
469 }
470
471 let increment = u64::try_from(tool_names.len()).unwrap_or(u64::MAX);
472 if let Some(limit) = self.limits.max_tool_calls {
473 let projected = self.tool_calls.saturating_add(increment);
474 if projected > limit {
475 return Some(count_preflight_exhaustion(
476 BudgetDimension::ToolCalls,
477 None,
478 limit,
479 self.tool_calls,
480 increment,
481 boundary,
482 ));
483 }
484 }
485
486 let mut batch_by_name = BTreeMap::<String, u64>::new();
487 for name in tool_names {
488 *batch_by_name.entry(name.clone()).or_default() += 1;
489 }
490 for (name, limit) in &self.limits.max_tool_calls_by_name {
491 let increment = batch_by_name.get(name).copied().unwrap_or_default();
492 if increment == 0 {
493 continue;
494 }
495 let observed = self
496 .tool_calls_by_name
497 .get(name)
498 .copied()
499 .unwrap_or_default();
500 if observed.saturating_add(increment) > *limit {
501 return Some(count_preflight_exhaustion(
502 BudgetDimension::ToolCallsByName,
503 Some(name.clone()),
504 *limit,
505 observed,
506 increment,
507 boundary,
508 ));
509 }
510 }
511
512 self.tool_calls = self.tool_calls.saturating_add(increment);
513 for (name, increment) in batch_by_name {
514 let value = self.tool_calls_by_name.entry(name).or_default();
515 *value = value.saturating_add(increment);
516 }
517 None
518 }
519
520 pub fn tool_batch_complete(&mut self, operation_failed: bool) -> Option<BudgetExhaustion> {
521 let boundary = BudgetEnforcementBoundary::ToolBatchComplete;
522 self.observe_boundary();
523 if operation_failed {
524 return None;
525 }
526 self.strict_unavailable(boundary).or_else(|| {
527 self.check_exceeded_limits(
528 boundary,
529 &[BudgetDimension::WallTime, BudgetDimension::HostCost],
530 )
531 })
532 }
533
534 pub fn terminal(&mut self) -> Option<BudgetExhaustion> {
535 let boundary = BudgetEnforcementBoundary::Terminal;
536 self.observe_boundary();
537 self.strict_unavailable(boundary)
538 .or_else(|| self.check_exceeded_limits(boundary, &DIMENSION_PRECEDENCE))
539 }
540
541 fn observe_boundary(&mut self) {
542 self.observe_elapsed();
543 self.observe_host_cost();
544 }
545
546 fn observe_elapsed(&mut self) {
547 if self.unavailable.contains_key(&BudgetDimension::WallTime) {
548 return;
549 }
550 let now_ns = (self.clock_ns)();
551 let delta_ns = now_ns.saturating_sub(self.started_ns);
552 let elapsed_ms = u128::from(self.base_elapsed_ms).saturating_add(delta_ns / 1_000_000);
553 if elapsed_ms > u128::from(MAX_WIRE_INTEGER) {
554 self.latch_unavailable(
555 BudgetDimension::WallTime,
556 BudgetUnavailableReason::IntegerOverflow,
557 Some("milliseconds"),
558 None,
559 None,
560 None,
561 );
562 return;
563 }
564 self.elapsed_ms = self.elapsed_ms.max(elapsed_ms as u64);
565 }
566
567 fn observe_host_cost(&mut self) {
568 let Some(limit) = self.limits.max_host_cost.clone() else {
569 return;
570 };
571 if self.unavailable.contains_key(&BudgetDimension::HostCost) {
572 return;
573 }
574 let Some(meter) = &self.host_cost_meter else {
575 self.latch_host_unavailable(&limit, BudgetUnavailableReason::MeterMissing, None);
576 return;
577 };
578 let reading = match meter.read() {
579 Ok(Some(reading)) => reading,
580 Ok(None) => {
581 self.latch_host_unavailable(
582 &limit,
583 BudgetUnavailableReason::MeterUnavailable,
584 None,
585 );
586 return;
587 }
588 Err(_) => {
589 self.latch_host_unavailable(&limit, BudgetUnavailableReason::MeterError, None);
590 return;
591 }
592 };
593 if reading.validate().is_err() {
594 self.latch_host_unavailable(&limit, BudgetUnavailableReason::MeterError, None);
595 return;
596 }
597 if reading.unit != limit.unit {
598 self.latch_host_unavailable(
599 &limit,
600 BudgetUnavailableReason::UnitMismatch,
601 Some(&reading),
602 );
603 return;
604 }
605 if reading.currency != limit.currency {
606 self.latch_host_unavailable(
607 &limit,
608 BudgetUnavailableReason::CurrencyMismatch,
609 Some(&reading),
610 );
611 return;
612 }
613 if self
614 .host_cost
615 .as_ref()
616 .is_some_and(|previous| reading.amount_microunits < previous.amount_microunits)
617 {
618 self.host_cost = None;
619 self.latch_host_unavailable(
620 &limit,
621 BudgetUnavailableReason::NonMonotonic,
622 Some(&reading),
623 );
624 return;
625 }
626 self.host_cost = Some(reading);
627 }
628
629 fn observe_token_usage(&mut self, usage: &TokenUsage) {
630 if !matches!(
631 usage.usage_source,
632 UsageSource::ProviderReported | UsageSource::Estimated
633 ) {
634 self.total_tokens = None;
635 self.latch_unavailable(
636 BudgetDimension::TotalTokens,
637 BudgetUnavailableReason::UsageMissing,
638 Some("tokens"),
639 None,
640 None,
641 None,
642 );
643 } else if let Some(current) = self.total_tokens {
644 self.total_tokens = self.safe_add_or_latch(
645 BudgetDimension::TotalTokens,
646 current,
647 usage.total_tokens,
648 "tokens",
649 );
650 }
651
652 let uncached = if usage.cache_usage.status == CacheUsageStatus::ProviderReported {
653 usage.cache_usage.uncached_input_tokens
654 } else {
655 None
656 };
657 if let Some(increment) = uncached {
658 if let Some(current) = self.uncached_input_tokens {
659 self.uncached_input_tokens = self.safe_add_or_latch(
660 BudgetDimension::UncachedInputTokens,
661 current,
662 increment,
663 "tokens",
664 );
665 }
666 } else {
667 self.uncached_input_tokens = None;
668 self.latch_unavailable(
669 BudgetDimension::UncachedInputTokens,
670 BudgetUnavailableReason::UsageMissing,
671 Some("tokens"),
672 None,
673 None,
674 None,
675 );
676 }
677 }
678
679 fn safe_add_or_latch(
680 &mut self,
681 dimension: BudgetDimension,
682 current: u64,
683 increment: u64,
684 expected_unit: &str,
685 ) -> Option<u64> {
686 match current.checked_add(increment) {
687 Some(total) if total <= MAX_WIRE_INTEGER => Some(total),
688 _ => {
689 self.latch_unavailable(
690 dimension,
691 BudgetUnavailableReason::IntegerOverflow,
692 Some(expected_unit),
693 None,
694 None,
695 None,
696 );
697 None
698 }
699 }
700 }
701
702 fn latch_host_unavailable(
703 &mut self,
704 limit: &HostCost,
705 reason: BudgetUnavailableReason,
706 reading: Option<&HostCost>,
707 ) {
708 self.latch_unavailable(
709 BudgetDimension::HostCost,
710 reason,
711 Some(&limit.unit),
712 reading.map(|value| value.unit.as_str()),
713 limit.currency.as_deref(),
714 reading.and_then(|value| value.currency.as_deref()),
715 );
716 }
717
718 #[allow(clippy::too_many_arguments)]
719 fn latch_unavailable(
720 &mut self,
721 dimension: BudgetDimension,
722 reason: BudgetUnavailableReason,
723 expected_unit: Option<&str>,
724 observed_unit: Option<&str>,
725 expected_currency: Option<&str>,
726 observed_currency: Option<&str>,
727 ) {
728 self.unavailable
729 .entry(dimension)
730 .or_insert_with(|| BudgetUnavailableDimension {
731 dimension,
732 reason,
733 expected_unit: expected_unit.map(str::to_string),
734 observed_unit: observed_unit.map(str::to_string),
735 expected_currency: expected_currency.map(str::to_string),
736 observed_currency: observed_currency.map(str::to_string),
737 });
738 }
739
740 fn strict_unavailable(&self, boundary: BudgetEnforcementBoundary) -> Option<BudgetExhaustion> {
741 if self.limits.unavailable_metric_policy != UnavailableMetricPolicy::Stop {
742 return None;
743 }
744 DIMENSION_PRECEDENCE.iter().find_map(|dimension| {
745 let unavailable = self.unavailable.get(dimension)?;
746 let (limit, unit, currency) = self.limit_descriptor(*dimension)?;
747 Some(BudgetExhaustion {
748 dimension: *dimension,
749 tool_name: None,
750 reason: BudgetExhaustionReason::MetricUnavailable,
751 limit,
752 observed: None,
753 attempted_increment: None,
754 overshoot: None,
755 unit,
756 currency,
757 enforcement_boundary: boundary,
758 unavailable_reason: Some(unavailable.reason),
759 })
760 })
761 }
762
763 fn check_admission_limits(
764 &self,
765 boundary: BudgetEnforcementBoundary,
766 dimensions: &[BudgetDimension],
767 ) -> Option<BudgetExhaustion> {
768 DIMENSION_PRECEDENCE.iter().find_map(|dimension| {
769 if !dimensions.contains(dimension) || self.unavailable.contains_key(dimension) {
770 return None;
771 }
772 let (limit, unit, currency) = self.limit_descriptor(*dimension)?;
773 let observed = self.observed_value(*dimension)?;
774 (observed >= limit).then(|| BudgetExhaustion {
775 dimension: *dimension,
776 tool_name: None,
777 reason: BudgetExhaustionReason::LimitReached,
778 limit,
779 observed: Some(observed),
780 attempted_increment: None,
781 overshoot: Some(observed.saturating_sub(limit)),
782 unit,
783 currency,
784 enforcement_boundary: boundary,
785 unavailable_reason: None,
786 })
787 })
788 }
789
790 fn check_exceeded_limits(
791 &self,
792 boundary: BudgetEnforcementBoundary,
793 dimensions: &[BudgetDimension],
794 ) -> Option<BudgetExhaustion> {
795 DIMENSION_PRECEDENCE.iter().find_map(|dimension| {
796 if !dimensions.contains(dimension) || self.unavailable.contains_key(dimension) {
797 return None;
798 }
799 let (limit, unit, currency) = self.limit_descriptor(*dimension)?;
800 let observed = self.observed_value(*dimension)?;
801 (observed > limit).then(|| BudgetExhaustion {
802 dimension: *dimension,
803 tool_name: None,
804 reason: BudgetExhaustionReason::LimitExceeded,
805 limit,
806 observed: Some(observed),
807 attempted_increment: None,
808 overshoot: Some(observed - limit),
809 unit,
810 currency,
811 enforcement_boundary: boundary,
812 unavailable_reason: None,
813 })
814 })
815 }
816
817 fn limit_descriptor(
818 &self,
819 dimension: BudgetDimension,
820 ) -> Option<(u64, String, Option<String>)> {
821 match dimension {
822 BudgetDimension::TotalTokens => self
823 .limits
824 .max_total_tokens
825 .map(|limit| (limit, "tokens".to_string(), None)),
826 BudgetDimension::UncachedInputTokens => self
827 .limits
828 .max_uncached_input_tokens
829 .map(|limit| (limit, "tokens".to_string(), None)),
830 BudgetDimension::ToolCalls => self
831 .limits
832 .max_tool_calls
833 .map(|limit| (limit, "calls".to_string(), None)),
834 BudgetDimension::WallTime => self
835 .limits
836 .max_wall_time_ms
837 .map(|limit| (limit, "milliseconds".to_string(), None)),
838 BudgetDimension::HostCost => self.limits.max_host_cost.as_ref().map(|cost| {
839 (
840 cost.amount_microunits,
841 cost.unit.clone(),
842 cost.currency.clone(),
843 )
844 }),
845 BudgetDimension::ToolCallsByName => None,
846 }
847 }
848
849 fn observed_value(&self, dimension: BudgetDimension) -> Option<u64> {
850 match dimension {
851 BudgetDimension::TotalTokens => self.total_tokens,
852 BudgetDimension::UncachedInputTokens => self.uncached_input_tokens,
853 BudgetDimension::ToolCalls => Some(self.tool_calls),
854 BudgetDimension::WallTime => Some(self.elapsed_ms),
855 BudgetDimension::HostCost => self.host_cost.as_ref().map(|cost| cost.amount_microunits),
856 BudgetDimension::ToolCallsByName => None,
857 }
858 }
859}
860
861const DIMENSION_PRECEDENCE: [BudgetDimension; 6] = [
862 BudgetDimension::WallTime,
863 BudgetDimension::TotalTokens,
864 BudgetDimension::UncachedInputTokens,
865 BudgetDimension::HostCost,
866 BudgetDimension::ToolCalls,
867 BudgetDimension::ToolCallsByName,
868];
869
870fn dimension_precedence(dimension: BudgetDimension) -> usize {
871 DIMENSION_PRECEDENCE
872 .iter()
873 .position(|candidate| *candidate == dimension)
874 .expect("all budget dimensions have stable precedence")
875}
876
877fn count_preflight_exhaustion(
878 dimension: BudgetDimension,
879 tool_name: Option<String>,
880 limit: u64,
881 observed: u64,
882 attempted_increment: u64,
883 boundary: BudgetEnforcementBoundary,
884) -> BudgetExhaustion {
885 BudgetExhaustion {
886 dimension,
887 tool_name,
888 reason: BudgetExhaustionReason::LimitReached,
889 limit,
890 observed: Some(observed),
891 attempted_increment: Some(attempted_increment),
892 overshoot: Some(
893 observed
894 .saturating_add(attempted_increment)
895 .saturating_sub(limit),
896 ),
897 unit: "calls".to_string(),
898 currency: None,
899 enforcement_boundary: boundary,
900 unavailable_reason: None,
901 }
902}
903
904fn validate_wire_integer(value: u64, field: &str) -> Result<(), String> {
905 if value <= MAX_WIRE_INTEGER {
906 Ok(())
907 } else {
908 Err(format!("{field} must be between 0 and {MAX_WIRE_INTEGER}"))
909 }
910}
911
912fn validate_non_empty(value: &str, field: &str) -> Result<(), String> {
913 if value.trim().is_empty() {
914 Err(format!("{field} must be a non-empty string"))
915 } else {
916 Ok(())
917 }
918}