Skip to main content

starweaver_usage/
lib.rs

1//! Usage accounting, limits, and optional pricing primitives for Starweaver.
2
3use 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/// Token and request usage accumulated by model and runtime layers.
17#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
18pub struct Usage {
19    /// Number of provider requests.
20    pub requests: u64,
21    /// Input or prompt tokens.
22    ///
23    /// Provider adapters should normalize this as total provider-billed input
24    /// tokens for the request, including cache-write and cache-read tokens when
25    /// those subtotals are present. Pricing helpers subtract the cache subtotals
26    /// before applying cache-specific rates.
27    pub input_tokens: u64,
28    /// Provider-reported total tokens written to a prompt cache across all durations.
29    ///
30    /// Use [`Self::effective_cache_write_tokens`] for derived accounting because
31    /// externally constructed or deserialized usage may contain inconsistent counters.
32    #[serde(default)]
33    pub cache_write_tokens: u64,
34    /// Provider-reported one-hour cache writes included in [`Self::cache_write_tokens`].
35    ///
36    /// Providers without a one-hour cache-write breakdown leave this at zero. Use
37    /// [`Self::effective_cache_write_1h_tokens`] for derived accounting.
38    #[serde(default, skip_serializing_if = "is_zero_u64")]
39    pub cache_write_1h_tokens: u64,
40    /// Tokens read from a provider prompt cache.
41    #[serde(default)]
42    pub cache_read_tokens: u64,
43    /// Output or completion tokens.
44    pub output_tokens: u64,
45    /// Total tokens.
46    pub total_tokens: u64,
47    /// Number of successful function tool calls executed by the runtime.
48    #[serde(default)]
49    pub tool_calls: u64,
50}
51
52impl Usage {
53    /// Return the effective aggregate cache-write tokens for derived accounting.
54    ///
55    /// Provider counters remain raw evidence. This projection repairs a one-hour
56    /// subset larger than the reported aggregate, then caps writes at the inclusive
57    /// input total so malformed counters cannot create additional billed tokens.
58    #[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    /// Return effective one-hour cache-write tokens for derived accounting.
73    #[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    /// Return effective cache-read tokens for derived accounting.
84    ///
85    /// Cache writes take precedence when malformed write and read subtotals exceed
86    /// inclusive input. Reads are capped at the input remaining after effective writes.
87    #[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    /// Return input tokens not assigned to an effective cache-write or read category.
100    #[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    /// Add another usage value into this one.
108    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    /// Return whether no usage has been recorded.
126    #[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    /// Return a copy with additional successful tool calls applied.
139    #[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/// Estimated USD pricing for usage.
147#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Ord, PartialOrd, Serialize)]
148pub struct PricingEstimate {
149    /// Estimated cost in micro USD units.
150    #[serde(default)]
151    pub amount_micros_usd: u64,
152}
153
154impl PricingEstimate {
155    /// Create an estimate from micro USD units.
156    #[must_use]
157    pub const fn from_micros_usd(amount_micros_usd: u64) -> Self {
158        Self { amount_micros_usd }
159    }
160
161    /// Add another estimate into this one.
162    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    /// Return whether the estimate is zero.
169    #[must_use]
170    pub const fn is_zero(&self) -> bool {
171        self.amount_micros_usd == 0
172    }
173}
174
175/// Cumulative usage for one agent or usage source in the current run.
176#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
177pub struct UsageSnapshotEntry {
178    /// Agent or source instance that generated this usage.
179    pub agent_id: String,
180    /// Human-readable agent or source name.
181    pub agent_name: String,
182    /// Model identifier that generated this usage.
183    pub model_id: String,
184    /// Cumulative token usage for this agent/source in the run.
185    pub usage: Usage,
186    /// Estimated cumulative pricing for this entry, in USD.
187    #[serde(default, skip_serializing_if = "Option::is_none")]
188    pub estimate_pricing: Option<PricingEstimate>,
189    /// Stable usage record id for idempotent updates.
190    #[serde(default, skip_serializing_if = "Option::is_none")]
191    pub usage_id: Option<String>,
192    /// Component that reported this usage.
193    #[serde(default = "default_usage_source")]
194    pub source: String,
195}
196
197/// Cumulative usage grouped by agent/source.
198#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
199pub struct UsageAgentTotal {
200    /// Human-readable agent or source name.
201    pub agent_name: String,
202    /// Model identifier, or `multiple` when a source used more than one model.
203    pub model_id: String,
204    /// Cumulative token usage for this agent/source.
205    pub usage: Usage,
206    /// Estimated cumulative pricing for this agent/source, in USD.
207    #[serde(default, skip_serializing_if = "Option::is_none")]
208    pub estimate_pricing: Option<PricingEstimate>,
209    /// Stable usage record id when all grouped entries share one id.
210    #[serde(default, skip_serializing_if = "Option::is_none")]
211    pub usage_id: Option<String>,
212    /// Component that reported this usage.
213    #[serde(default = "default_usage_source")]
214    pub source: String,
215}
216
217/// Cumulative usage snapshot for one run.
218///
219/// Realtime consumers should treat each snapshot as a replacement for the
220/// previous snapshot with the same run id.
221#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
222pub struct UsageSnapshot {
223    /// Run identifier for the snapshot.
224    pub run_id: String,
225    /// Usage reported by the latest provider request that produced this snapshot.
226    ///
227    /// This is intentionally separate from `total_usage`: realtime UI surfaces may use
228    /// the latest request total tokens as the current context-window estimate,
229    /// while `total_usage` remains the cumulative run ledger.
230    #[serde(default, skip_serializing_if = "Option::is_none")]
231    pub latest_usage: Option<Usage>,
232    /// Cumulative usage across all known entries in this run.
233    #[serde(default)]
234    pub total_usage: Usage,
235    /// Estimated cumulative pricing across all known entries in this run, in USD.
236    #[serde(default, skip_serializing_if = "Option::is_none")]
237    pub estimate_pricing: Option<PricingEstimate>,
238    /// Per-agent/source cumulative usage entries.
239    #[serde(default, skip_serializing_if = "Vec::is_empty")]
240    pub entries: Vec<UsageSnapshotEntry>,
241    /// Cumulative usage grouped by agent id.
242    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
243    pub agent_usages: BTreeMap<String, UsageAgentTotal>,
244    /// Cumulative usage grouped by model id.
245    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
246    pub model_usages: BTreeMap<String, Usage>,
247    /// Estimated cumulative pricing grouped by model id, in USD.
248    #[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/// Runtime usage limits for one agent run.
257#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
258pub struct UsageLimits {
259    /// Maximum provider requests allowed in one run.
260    #[serde(default, skip_serializing_if = "Option::is_none")]
261    pub request_limit: Option<u64>,
262    /// Maximum input tokens allowed.
263    #[serde(default, skip_serializing_if = "Option::is_none")]
264    pub input_tokens_limit: Option<u64>,
265    /// Maximum output tokens allowed.
266    #[serde(default, skip_serializing_if = "Option::is_none")]
267    pub output_tokens_limit: Option<u64>,
268    /// Maximum total tokens allowed.
269    #[serde(default, skip_serializing_if = "Option::is_none")]
270    pub total_tokens_limit: Option<u64>,
271    /// Maximum successful function tool calls allowed.
272    #[serde(default, skip_serializing_if = "Option::is_none")]
273    pub tool_calls_limit: Option<u64>,
274    /// Optional USD pricing budget based on accumulated usage.
275    #[cfg(feature = "pricing")]
276    #[serde(default, skip_serializing_if = "Option::is_none")]
277    pub cost_budget: Option<pricing::CostBudget>,
278}
279
280impl UsageLimits {
281    /// Create empty limits.
282    #[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    /// Set request limit.
296    #[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    /// Set input token limit.
303    #[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    /// Set output token limit.
310    #[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    /// Set total token limit.
317    #[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    /// Set successful tool-call limit.
324    #[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    /// Set USD cost budget.
331    #[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    /// Estimate current USD cost in micro-units when a cost budget is configured.
339    #[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    /// Estimate current USD pricing when a cost budget is configured.
348    #[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    /// Check whether the next model request would exceed the request limit.
357    ///
358    /// # Errors
359    ///
360    /// Returns an error when another request would exceed the configured request limit.
361    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    /// Check whether projected tool calls would exceed the tool-call limit.
375    ///
376    /// # Errors
377    ///
378    /// Returns an error when executing the next successful tool calls would exceed the configured limit.
379    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    /// Check whether accumulated usage exceeds configured token or pricing limits.
390    ///
391    /// # Errors
392    ///
393    /// Returns an error when accumulated usage exceeds any configured limit.
394    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/// Token counter checked by a usage limit.
419#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
420#[serde(rename_all = "snake_case")]
421pub enum UsageTokenKind {
422    /// Input or prompt tokens.
423    InputTokens,
424    /// Output or completion tokens.
425    OutputTokens,
426    /// Total tokens.
427    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/// Usage limit error.
442#[derive(Clone, Debug, Error, Deserialize, Eq, PartialEq, Serialize)]
443pub enum UsageLimitError {
444    /// The next request would exceed request budget.
445    #[error(
446        "the next request would exceed the request_limit of {limit} (next_requests={next_requests})"
447    )]
448    NextRequest {
449        /// Configured limit.
450        limit: u64,
451        /// Requests after the next request.
452        next_requests: u64,
453    },
454    /// Accumulated usage exceeded a token budget.
455    #[error("exceeded the {kind}_limit of {limit} ({kind}={actual})")]
456    Token {
457        /// Usage kind.
458        kind: UsageTokenKind,
459        /// Configured limit.
460        limit: u64,
461        /// Actual usage value.
462        actual: u64,
463    },
464    /// Accumulated usage exceeded a USD pricing budget.
465    #[cfg(feature = "pricing")]
466    #[error("exceeded the total_cost_limit_micros of {limit_micros} (cost_micros={actual_micros})")]
467    Cost {
468        /// Configured cost limit in micro USD units.
469        limit_micros: u64,
470        /// Actual estimated cost in micro USD units.
471        actual_micros: u64,
472    },
473    /// Projected successful function tool calls would exceed the configured budget.
474    #[error(
475        "the next tool call(s) would exceed the tool_calls_limit of {limit} (tool_calls={tool_calls})"
476    )]
477    ToolCalls {
478        /// Configured tool-call limit.
479        limit: u64,
480        /// Projected successful tool calls.
481        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
500/// Aggregate an optional pricing estimate into a running total.
501pub 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}