1use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6use thiserror::Error;
7
8#[allow(clippy::trivially_copy_pass_by_ref)]
9const fn is_zero_u64(value: &u64) -> bool {
10 *value == 0
11}
12
13#[cfg(feature = "pricing")]
14pub mod pricing;
15
16#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
18pub struct Usage {
19 pub requests: u64,
21 pub input_tokens: u64,
28 #[serde(default)]
33 pub cache_write_tokens: u64,
34 #[serde(default, skip_serializing_if = "is_zero_u64")]
39 pub cache_write_1h_tokens: u64,
40 #[serde(default)]
42 pub cache_read_tokens: u64,
43 pub output_tokens: u64,
45 pub total_tokens: u64,
47 #[serde(default)]
49 pub tool_calls: u64,
50}
51
52impl Usage {
53 #[must_use]
59 pub const fn effective_cache_write_tokens(&self) -> u64 {
60 let reported = if self.cache_write_tokens > self.cache_write_1h_tokens {
61 self.cache_write_tokens
62 } else {
63 self.cache_write_1h_tokens
64 };
65 if reported < self.input_tokens {
66 reported
67 } else {
68 self.input_tokens
69 }
70 }
71
72 #[must_use]
74 pub const fn effective_cache_write_1h_tokens(&self) -> u64 {
75 let cache_write_tokens = self.effective_cache_write_tokens();
76 if self.cache_write_1h_tokens < cache_write_tokens {
77 self.cache_write_1h_tokens
78 } else {
79 cache_write_tokens
80 }
81 }
82
83 #[must_use]
88 pub const fn effective_cache_read_tokens(&self) -> u64 {
89 let remaining_input = self
90 .input_tokens
91 .saturating_sub(self.effective_cache_write_tokens());
92 if self.cache_read_tokens < remaining_input {
93 self.cache_read_tokens
94 } else {
95 remaining_input
96 }
97 }
98
99 #[must_use]
101 pub const fn effective_standard_input_tokens(&self) -> u64 {
102 self.input_tokens
103 .saturating_sub(self.effective_cache_write_tokens())
104 .saturating_sub(self.effective_cache_read_tokens())
105 }
106
107 pub const fn add_assign(&mut self, other: &Self) {
109 self.requests = self.requests.saturating_add(other.requests);
110 self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
111 self.cache_write_tokens = self
112 .cache_write_tokens
113 .saturating_add(other.cache_write_tokens);
114 self.cache_write_1h_tokens = self
115 .cache_write_1h_tokens
116 .saturating_add(other.cache_write_1h_tokens);
117 self.cache_read_tokens = self
118 .cache_read_tokens
119 .saturating_add(other.cache_read_tokens);
120 self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
121 self.total_tokens = self.total_tokens.saturating_add(other.total_tokens);
122 self.tool_calls = self.tool_calls.saturating_add(other.tool_calls);
123 }
124
125 #[must_use]
127 pub const fn is_empty(&self) -> bool {
128 self.requests == 0
129 && self.input_tokens == 0
130 && self.cache_write_tokens == 0
131 && self.cache_write_1h_tokens == 0
132 && self.cache_read_tokens == 0
133 && self.output_tokens == 0
134 && self.total_tokens == 0
135 && self.tool_calls == 0
136 }
137
138 #[must_use]
140 pub const fn with_additional_tool_calls(mut self, tool_calls: u64) -> Self {
141 self.tool_calls = self.tool_calls.saturating_add(tool_calls);
142 self
143 }
144}
145
146#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Ord, PartialOrd, Serialize)]
148pub struct PricingEstimate {
149 #[serde(default)]
151 pub amount_micros_usd: u64,
152}
153
154impl PricingEstimate {
155 #[must_use]
157 pub const fn from_micros_usd(amount_micros_usd: u64) -> Self {
158 Self { amount_micros_usd }
159 }
160
161 pub const fn add_assign(&mut self, other: &Self) {
163 self.amount_micros_usd = self
164 .amount_micros_usd
165 .saturating_add(other.amount_micros_usd);
166 }
167
168 #[must_use]
170 pub const fn is_zero(&self) -> bool {
171 self.amount_micros_usd == 0
172 }
173}
174
175#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
177pub struct UsageSnapshotEntry {
178 pub agent_id: String,
180 pub agent_name: String,
182 pub model_id: String,
184 pub usage: Usage,
186 #[serde(default, skip_serializing_if = "Option::is_none")]
188 pub estimate_pricing: Option<PricingEstimate>,
189 #[serde(default, skip_serializing_if = "Option::is_none")]
191 pub usage_id: Option<String>,
192 #[serde(default = "default_usage_source")]
194 pub source: String,
195}
196
197#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
199pub struct UsageAgentTotal {
200 pub agent_name: String,
202 pub model_id: String,
204 pub usage: Usage,
206 #[serde(default, skip_serializing_if = "Option::is_none")]
208 pub estimate_pricing: Option<PricingEstimate>,
209 #[serde(default, skip_serializing_if = "Option::is_none")]
211 pub usage_id: Option<String>,
212 #[serde(default = "default_usage_source")]
214 pub source: String,
215}
216
217#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
222pub struct UsageSnapshot {
223 pub run_id: String,
225 #[serde(default, skip_serializing_if = "Option::is_none")]
231 pub latest_usage: Option<Usage>,
232 #[serde(default)]
234 pub total_usage: Usage,
235 #[serde(default, skip_serializing_if = "Option::is_none")]
237 pub estimate_pricing: Option<PricingEstimate>,
238 #[serde(default, skip_serializing_if = "Vec::is_empty")]
240 pub entries: Vec<UsageSnapshotEntry>,
241 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
243 pub agent_usages: BTreeMap<String, UsageAgentTotal>,
244 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
246 pub model_usages: BTreeMap<String, Usage>,
247 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
249 pub model_estimate_pricing: BTreeMap<String, PricingEstimate>,
250}
251
252fn default_usage_source() -> String {
253 "model_request".to_string()
254}
255
256#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
258pub struct UsageLimits {
259 #[serde(default, skip_serializing_if = "Option::is_none")]
261 pub request_limit: Option<u64>,
262 #[serde(default, skip_serializing_if = "Option::is_none")]
264 pub input_tokens_limit: Option<u64>,
265 #[serde(default, skip_serializing_if = "Option::is_none")]
267 pub output_tokens_limit: Option<u64>,
268 #[serde(default, skip_serializing_if = "Option::is_none")]
270 pub total_tokens_limit: Option<u64>,
271 #[serde(default, skip_serializing_if = "Option::is_none")]
273 pub tool_calls_limit: Option<u64>,
274 #[cfg(feature = "pricing")]
276 #[serde(default, skip_serializing_if = "Option::is_none")]
277 pub cost_budget: Option<pricing::CostBudget>,
278}
279
280impl UsageLimits {
281 #[must_use]
283 pub const fn new() -> Self {
284 Self {
285 request_limit: None,
286 input_tokens_limit: None,
287 output_tokens_limit: None,
288 total_tokens_limit: None,
289 tool_calls_limit: None,
290 #[cfg(feature = "pricing")]
291 cost_budget: None,
292 }
293 }
294
295 #[must_use]
297 pub const fn with_request_limit(mut self, limit: u64) -> Self {
298 self.request_limit = Some(limit);
299 self
300 }
301
302 #[must_use]
304 pub const fn with_input_tokens_limit(mut self, limit: u64) -> Self {
305 self.input_tokens_limit = Some(limit);
306 self
307 }
308
309 #[must_use]
311 pub const fn with_output_tokens_limit(mut self, limit: u64) -> Self {
312 self.output_tokens_limit = Some(limit);
313 self
314 }
315
316 #[must_use]
318 pub const fn with_total_tokens_limit(mut self, limit: u64) -> Self {
319 self.total_tokens_limit = Some(limit);
320 self
321 }
322
323 #[must_use]
325 pub const fn with_tool_calls_limit(mut self, limit: u64) -> Self {
326 self.tool_calls_limit = Some(limit);
327 self
328 }
329
330 #[cfg(feature = "pricing")]
332 #[must_use]
333 pub const fn with_cost_budget(mut self, budget: pricing::CostBudget) -> Self {
334 self.cost_budget = Some(budget);
335 self
336 }
337
338 #[cfg(feature = "pricing")]
340 #[must_use]
341 pub fn estimate_cost_micros(&self, usage: &Usage) -> Option<u64> {
342 self.cost_budget
343 .as_ref()
344 .map(|budget| budget.estimate_micros(usage))
345 }
346
347 #[cfg(feature = "pricing")]
349 #[must_use]
350 pub fn estimate_pricing(&self, usage: &Usage) -> Option<PricingEstimate> {
351 self.cost_budget
352 .as_ref()
353 .map(|budget| budget.estimate_pricing(usage))
354 }
355
356 pub const fn check_before_request(&self, current: &Usage) -> Result<(), UsageLimitError> {
362 if let Some(limit) = self.request_limit {
363 let next = current.requests.saturating_add(1);
364 if next > limit {
365 return Err(UsageLimitError::NextRequest {
366 limit,
367 next_requests: next,
368 });
369 }
370 }
371 Ok(())
372 }
373
374 pub const fn check_tool_calls(&self, projected: &Usage) -> Result<(), UsageLimitError> {
380 match self.tool_calls_limit {
381 Some(limit) if projected.tool_calls > limit => Err(UsageLimitError::ToolCalls {
382 limit,
383 tool_calls: projected.tool_calls,
384 }),
385 _ => Ok(()),
386 }
387 }
388
389 pub fn check_usage(&self, usage: &Usage) -> Result<(), UsageLimitError> {
395 check_limit(
396 UsageTokenKind::InputTokens,
397 self.input_tokens_limit,
398 usage.input_tokens,
399 )?;
400 check_limit(
401 UsageTokenKind::OutputTokens,
402 self.output_tokens_limit,
403 usage.output_tokens,
404 )?;
405 check_limit(
406 UsageTokenKind::TotalTokens,
407 self.total_tokens_limit,
408 usage.total_tokens,
409 )?;
410 #[cfg(feature = "pricing")]
411 if let Some(budget) = &self.cost_budget {
412 budget.check_usage(usage)?;
413 }
414 Ok(())
415 }
416}
417
418#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
420#[serde(rename_all = "snake_case")]
421pub enum UsageTokenKind {
422 InputTokens,
424 OutputTokens,
426 TotalTokens,
428}
429
430impl std::fmt::Display for UsageTokenKind {
431 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
432 let value = match self {
433 Self::InputTokens => "input_tokens",
434 Self::OutputTokens => "output_tokens",
435 Self::TotalTokens => "total_tokens",
436 };
437 formatter.write_str(value)
438 }
439}
440
441#[derive(Clone, Debug, Error, Deserialize, Eq, PartialEq, Serialize)]
443pub enum UsageLimitError {
444 #[error(
446 "the next request would exceed the request_limit of {limit} (next_requests={next_requests})"
447 )]
448 NextRequest {
449 limit: u64,
451 next_requests: u64,
453 },
454 #[error("exceeded the {kind}_limit of {limit} ({kind}={actual})")]
456 Token {
457 kind: UsageTokenKind,
459 limit: u64,
461 actual: u64,
463 },
464 #[cfg(feature = "pricing")]
466 #[error("exceeded the total_cost_limit_micros of {limit_micros} (cost_micros={actual_micros})")]
467 Cost {
468 limit_micros: u64,
470 actual_micros: u64,
472 },
473 #[error(
475 "the next tool call(s) would exceed the tool_calls_limit of {limit} (tool_calls={tool_calls})"
476 )]
477 ToolCalls {
478 limit: u64,
480 tool_calls: u64,
482 },
483}
484
485const fn check_limit(
486 kind: UsageTokenKind,
487 limit: Option<u64>,
488 actual: u64,
489) -> Result<(), UsageLimitError> {
490 match limit {
491 Some(limit) if actual > limit => Err(UsageLimitError::Token {
492 kind,
493 limit,
494 actual,
495 }),
496 _ => Ok(()),
497 }
498}
499
500pub fn add_optional_pricing(
502 total: &mut Option<PricingEstimate>,
503 estimate: Option<&PricingEstimate>,
504) {
505 if let Some(estimate) = estimate {
506 match total {
507 Some(total) => total.add_assign(estimate),
508 None => *total = Some(estimate.clone()),
509 }
510 }
511}
512
513#[cfg(test)]
514mod tests {
515 use super::*;
516
517 #[test]
518 fn usage_add_assign_and_empty_work() {
519 let mut usage = Usage {
520 requests: 1,
521 input_tokens: 2,
522 cache_write_tokens: 7,
523 cache_write_1h_tokens: 3,
524 cache_read_tokens: 11,
525 output_tokens: 3,
526 total_tokens: 5,
527 tool_calls: 1,
528 };
529 usage.add_assign(&Usage {
530 requests: 2,
531 input_tokens: 4,
532 cache_write_tokens: 13,
533 cache_write_1h_tokens: 5,
534 cache_read_tokens: 17,
535 output_tokens: 6,
536 total_tokens: 10,
537 tool_calls: 3,
538 });
539 assert_eq!(usage.requests, 3);
540 assert_eq!(usage.input_tokens, 6);
541 assert_eq!(usage.cache_write_tokens, 20);
542 assert_eq!(usage.cache_write_1h_tokens, 8);
543 assert_eq!(usage.cache_read_tokens, 28);
544 assert_eq!(usage.output_tokens, 9);
545 assert_eq!(usage.total_tokens, 15);
546 assert_eq!(usage.tool_calls, 4);
547 assert_eq!(usage.clone().with_additional_tool_calls(2).tool_calls, 6);
548 assert!(Usage::default().is_empty());
549 assert!(!usage.is_empty());
550 }
551
552 #[test]
553 fn usage_add_assign_saturates() {
554 let mut usage = Usage {
555 requests: u64::MAX,
556 input_tokens: u64::MAX,
557 cache_write_tokens: u64::MAX,
558 cache_write_1h_tokens: u64::MAX,
559 cache_read_tokens: u64::MAX,
560 output_tokens: u64::MAX,
561 total_tokens: u64::MAX,
562 tool_calls: u64::MAX,
563 };
564 usage.add_assign(&Usage {
565 requests: 1,
566 input_tokens: 1,
567 cache_write_tokens: 1,
568 cache_write_1h_tokens: 1,
569 cache_read_tokens: 1,
570 output_tokens: 1,
571 total_tokens: 1,
572 tool_calls: 1,
573 });
574 assert_eq!(usage.requests, u64::MAX);
575 assert_eq!(usage.input_tokens, u64::MAX);
576 assert_eq!(usage.cache_write_tokens, u64::MAX);
577 assert_eq!(usage.cache_write_1h_tokens, u64::MAX);
578 assert_eq!(usage.cache_read_tokens, u64::MAX);
579 assert_eq!(usage.output_tokens, u64::MAX);
580 assert_eq!(usage.total_tokens, u64::MAX);
581 assert_eq!(usage.tool_calls, u64::MAX);
582 }
583
584 #[test]
585 fn effective_cache_breakdown_bounds_malformed_reported_counters() {
586 let subset_exceeds_aggregate = Usage {
587 input_tokens: 300,
588 cache_write_tokens: 100,
589 cache_write_1h_tokens: 200,
590 cache_read_tokens: 50,
591 ..Usage::default()
592 };
593 assert_eq!(subset_exceeds_aggregate.effective_cache_write_tokens(), 200);
594 assert_eq!(
595 subset_exceeds_aggregate.effective_cache_write_1h_tokens(),
596 200
597 );
598 assert_eq!(subset_exceeds_aggregate.effective_cache_read_tokens(), 50);
599 assert_eq!(
600 subset_exceeds_aggregate.effective_standard_input_tokens(),
601 50
602 );
603
604 let aggregate_exceeds_input = Usage {
605 input_tokens: 100,
606 cache_write_tokens: 200,
607 cache_write_1h_tokens: 50,
608 cache_read_tokens: 10,
609 ..Usage::default()
610 };
611 assert_eq!(aggregate_exceeds_input.effective_cache_write_tokens(), 100);
612 assert_eq!(
613 aggregate_exceeds_input.effective_cache_write_1h_tokens(),
614 50
615 );
616 assert_eq!(aggregate_exceeds_input.effective_cache_read_tokens(), 0);
617 assert_eq!(aggregate_exceeds_input.effective_standard_input_tokens(), 0);
618
619 let writes_and_reads_exceed_input = Usage {
620 input_tokens: 100,
621 cache_write_tokens: 60,
622 cache_write_1h_tokens: 20,
623 cache_read_tokens: 70,
624 ..Usage::default()
625 };
626 assert_eq!(
627 writes_and_reads_exceed_input.effective_cache_write_tokens(),
628 60
629 );
630 assert_eq!(
631 writes_and_reads_exceed_input.effective_cache_write_1h_tokens(),
632 20
633 );
634 assert_eq!(
635 writes_and_reads_exceed_input.effective_cache_read_tokens(),
636 40
637 );
638 assert_eq!(
639 writes_and_reads_exceed_input.effective_standard_input_tokens(),
640 0
641 );
642
643 let near_overflow = Usage {
644 input_tokens: u64::MAX,
645 cache_write_tokens: u64::MAX - 1,
646 cache_write_1h_tokens: u64::MAX,
647 cache_read_tokens: u64::MAX,
648 ..Usage::default()
649 };
650 assert_eq!(near_overflow.effective_cache_write_tokens(), u64::MAX);
651 assert_eq!(near_overflow.effective_cache_write_1h_tokens(), u64::MAX);
652 assert_eq!(near_overflow.effective_cache_read_tokens(), 0);
653 assert_eq!(near_overflow.effective_standard_input_tokens(), 0);
654 }
655
656 #[test]
657 fn usage_one_hour_cache_write_tokens_are_backward_compatible_in_json() {
658 let legacy = serde_json::json!({
659 "requests": 1,
660 "input_tokens": 2,
661 "cache_write_tokens": 3,
662 "cache_read_tokens": 4,
663 "output_tokens": 5,
664 "total_tokens": 11,
665 "tool_calls": 0
666 });
667 let usage = match serde_json::from_value::<Usage>(legacy) {
668 Ok(usage) => usage,
669 Err(err) => panic!("legacy usage should deserialize: {err}"),
670 };
671 assert_eq!(usage.cache_write_1h_tokens, 0);
672
673 let serialized = match serde_json::to_value(&usage) {
674 Ok(value) => value,
675 Err(err) => panic!("usage should serialize: {err}"),
676 };
677 assert!(serialized.get("cache_write_1h_tokens").is_none());
678
679 let one_hour = Usage {
680 cache_write_tokens: 3,
681 cache_write_1h_tokens: 2,
682 ..Usage::default()
683 };
684 assert_eq!(
685 serde_json::to_value(one_hour)
686 .ok()
687 .and_then(|value| value["cache_write_1h_tokens"].as_u64()),
688 Some(2)
689 );
690 }
691
692 #[test]
693 fn usage_limit_error_token_kind_is_owned_ser_de_contract() {
694 let error = UsageLimitError::Token {
695 kind: UsageTokenKind::TotalTokens,
696 limit: 5,
697 actual: 6,
698 };
699 let value = match serde_json::to_value(&error) {
700 Ok(value) => value,
701 Err(err) => panic!("usage limit error should serialize: {err}"),
702 };
703 let restored: UsageLimitError = match serde_json::from_value(value) {
704 Ok(restored) => restored,
705 Err(err) => panic!("usage limit error should deserialize: {err}"),
706 };
707 assert_eq!(restored, error);
708 }
709
710 #[test]
711 fn usage_snapshot_accepts_missing_pricing_fields() {
712 let snapshot: UsageSnapshot = match serde_json::from_value(serde_json::json!({
713 "run_id": "run_1",
714 "total_usage": {
715 "requests": 1,
716 "input_tokens": 2,
717 "output_tokens": 3,
718 "total_tokens": 5
719 }
720 })) {
721 Ok(snapshot) => snapshot,
722 Err(err) => panic!("usage snapshot should deserialize: {err}"),
723 };
724 assert_eq!(snapshot.run_id, "run_1");
725 assert!(snapshot.estimate_pricing.is_none());
726 assert!(snapshot.model_estimate_pricing.is_empty());
727 }
728}