tea_protocol/
usage.rs

1use std::fmt;
2use std::str::FromStr;
3
4use serde::{Deserialize, Deserializer, Serialize, Serializer};
5use thiserror::Error;
6
7/// Largest integer exactly representable by a JavaScript `Number`.
8pub const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991;
9const MAX_DECIMAL_DIGITS: usize = 36;
10const MAX_DECIMAL_SCALE: usize = 18;
11
12/// A token count safe to encode as a JSON number for JavaScript consumers.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
14#[serde(transparent)]
15pub struct TokenCount(u64);
16
17impl TokenCount {
18    /// Creates a JavaScript-safe token count.
19    ///
20    /// # Errors
21    ///
22    /// Returns [`UsageError::UnsafeInteger`] when `value` exceeds
23    /// [`MAX_SAFE_INTEGER`].
24    pub const fn new(value: u64) -> Result<Self, UsageError> {
25        if value <= MAX_SAFE_INTEGER {
26            Ok(Self(value))
27        } else {
28            Err(UsageError::UnsafeInteger)
29        }
30    }
31
32    /// Returns the integer token count.
33    #[must_use]
34    pub const fn get(self) -> u64 {
35        self.0
36    }
37}
38
39impl<'de> Deserialize<'de> for TokenCount {
40    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
41    where
42        D: Deserializer<'de>,
43    {
44        let value = u64::deserialize(deserializer)?;
45        Self::new(value).map_err(serde::de::Error::custom)
46    }
47}
48
49/// Provider-neutral token usage.
50#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
51#[serde(rename_all = "camelCase")]
52#[allow(clippy::struct_field_names)] // `_tokens` is stable domain and wire vocabulary.
53pub struct Usage {
54    input_tokens: TokenCount,
55    output_tokens: TokenCount,
56    #[serde(skip_serializing_if = "Option::is_none")]
57    cache_read_tokens: Option<TokenCount>,
58    #[serde(skip_serializing_if = "Option::is_none")]
59    cache_write_tokens: Option<TokenCount>,
60    #[serde(skip_serializing_if = "Option::is_none")]
61    reasoning_tokens: Option<TokenCount>,
62}
63
64impl Usage {
65    /// Creates usage with required input and output token counts.
66    #[must_use]
67    pub const fn new(input_tokens: TokenCount, output_tokens: TokenCount) -> Self {
68        Self {
69            input_tokens,
70            output_tokens,
71            cache_read_tokens: None,
72            cache_write_tokens: None,
73            reasoning_tokens: None,
74        }
75    }
76
77    /// Returns input tokens.
78    #[must_use]
79    pub const fn input_tokens(&self) -> TokenCount {
80        self.input_tokens
81    }
82
83    /// Returns output tokens.
84    #[must_use]
85    pub const fn output_tokens(&self) -> TokenCount {
86        self.output_tokens
87    }
88
89    /// Returns cache-read tokens when reported.
90    #[must_use]
91    pub const fn cache_read_tokens(&self) -> Option<TokenCount> {
92        self.cache_read_tokens
93    }
94
95    /// Returns cache-write tokens when reported.
96    #[must_use]
97    pub const fn cache_write_tokens(&self) -> Option<TokenCount> {
98        self.cache_write_tokens
99    }
100
101    /// Returns reasoning tokens when reported.
102    #[must_use]
103    pub const fn reasoning_tokens(&self) -> Option<TokenCount> {
104        self.reasoning_tokens
105    }
106
107    /// Adds the cache-read token count.
108    #[must_use]
109    pub const fn with_cache_read(mut self, value: TokenCount) -> Self {
110        self.cache_read_tokens = Some(value);
111        self
112    }
113
114    /// Adds the cache-write token count.
115    #[must_use]
116    pub const fn with_cache_write(mut self, value: TokenCount) -> Self {
117        self.cache_write_tokens = Some(value);
118        self
119    }
120
121    /// Adds reasoning tokens, which must be a subset of output tokens.
122    ///
123    /// # Errors
124    ///
125    /// Returns [`UsageError::ReasoningExceedsOutput`] when the reasoning count
126    /// exceeds the output count.
127    pub fn with_reasoning(mut self, value: TokenCount) -> Result<Self, UsageError> {
128        if value > self.output_tokens {
129            return Err(UsageError::ReasoningExceedsOutput);
130        }
131        self.reasoning_tokens = Some(value);
132        Ok(self)
133    }
134
135    /// Returns the total billable/context token count without double-counting reasoning.
136    ///
137    /// # Errors
138    ///
139    /// Returns [`UsageError::TotalOverflow`] when the sum overflows or exceeds
140    /// the JSON safe-integer range.
141    pub fn total_tokens(&self) -> Result<TokenCount, UsageError> {
142        let total = [
143            Some(self.input_tokens),
144            Some(self.output_tokens),
145            self.cache_read_tokens,
146            self.cache_write_tokens,
147        ]
148        .into_iter()
149        .flatten()
150        .try_fold(0_u64, |total, value| total.checked_add(value.get()))
151        .ok_or(UsageError::TotalOverflow)?;
152        TokenCount::new(total)
153    }
154}
155
156#[derive(Deserialize)]
157#[serde(rename_all = "camelCase")]
158struct RawUsage {
159    #[serde(rename = "inputTokens")]
160    input: TokenCount,
161    #[serde(rename = "outputTokens")]
162    output: TokenCount,
163    #[serde(default, rename = "cacheReadTokens")]
164    cache_read: Option<TokenCount>,
165    #[serde(default, rename = "cacheWriteTokens")]
166    cache_write: Option<TokenCount>,
167    #[serde(default, rename = "reasoningTokens")]
168    reasoning: Option<TokenCount>,
169}
170
171impl<'de> Deserialize<'de> for Usage {
172    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
173    where
174        D: Deserializer<'de>,
175    {
176        let raw = RawUsage::deserialize(deserializer)?;
177        let mut usage = Self::new(raw.input, raw.output);
178        usage.cache_read_tokens = raw.cache_read;
179        usage.cache_write_tokens = raw.cache_write;
180        if let Some(reasoning) = raw.reasoning {
181            usage = usage
182                .with_reasoning(reasoning)
183                .map_err(serde::de::Error::custom)?;
184        }
185        usage.total_tokens().map_err(serde::de::Error::custom)?;
186        Ok(usage)
187    }
188}
189
190/// Error returned when validating token usage.
191#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
192pub enum UsageError {
193    /// A count exceeds JavaScript's safe integer range.
194    #[error("token count exceeds the JSON safe-integer range")]
195    UnsafeInteger,
196    /// Reasoning tokens exceed output tokens.
197    #[error("reasoning tokens must be a subset of output tokens")]
198    ReasoningExceedsOutput,
199    /// Summing usage counts overflowed or exceeded the safe range.
200    #[error("total token count exceeds the supported range")]
201    TotalOverflow,
202}
203
204/// A canonical, non-negative decimal amount encoded as text.
205#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
206pub struct DecimalAmount(String);
207
208impl DecimalAmount {
209    /// Returns the canonical decimal representation.
210    #[must_use]
211    pub fn as_str(&self) -> &str {
212        &self.0
213    }
214}
215
216/// Error returned when parsing an exact decimal amount.
217#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
218pub enum DecimalAmountParseError {
219    /// The amount is not canonical non-negative decimal text.
220    #[error("amount must use canonical non-negative decimal text")]
221    InvalidFormat,
222    /// The amount exceeds the supported precision or scale.
223    #[error("amount exceeds the supported precision or scale")]
224    TooPrecise,
225}
226
227impl FromStr for DecimalAmount {
228    type Err = DecimalAmountParseError;
229
230    fn from_str(value: &str) -> Result<Self, Self::Err> {
231        let (integer, fraction) = value
232            .split_once('.')
233            .map_or((value, None), |(left, right)| (left, Some(right)));
234        if integer.is_empty()
235            || !integer.bytes().all(|byte| byte.is_ascii_digit())
236            || (integer.len() > 1 && integer.starts_with('0'))
237        {
238            return Err(DecimalAmountParseError::InvalidFormat);
239        }
240        if let Some(fraction) = fraction {
241            if fraction.is_empty()
242                || !fraction.bytes().all(|byte| byte.is_ascii_digit())
243                || fraction.ends_with('0')
244            {
245                return Err(DecimalAmountParseError::InvalidFormat);
246            }
247            if fraction.len() > MAX_DECIMAL_SCALE {
248                return Err(DecimalAmountParseError::TooPrecise);
249            }
250        }
251        let digits = integer.len() + fraction.map_or(0, str::len);
252        if digits > MAX_DECIMAL_DIGITS {
253            return Err(DecimalAmountParseError::TooPrecise);
254        }
255        Ok(Self(value.to_owned()))
256    }
257}
258
259impl fmt::Display for DecimalAmount {
260    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
261        formatter.write_str(&self.0)
262    }
263}
264
265impl Serialize for DecimalAmount {
266    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
267    where
268        S: Serializer,
269    {
270        serializer.serialize_str(&self.0)
271    }
272}
273
274impl<'de> Deserialize<'de> for DecimalAmount {
275    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
276    where
277        D: Deserializer<'de>,
278    {
279        String::deserialize(deserializer)?
280            .parse()
281            .map_err(serde::de::Error::custom)
282    }
283}
284
285/// A three-letter uppercase ISO-style currency code.
286#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
287pub struct CurrencyCode(String);
288
289impl CurrencyCode {
290    /// Returns the three-letter code.
291    #[must_use]
292    pub fn as_str(&self) -> &str {
293        &self.0
294    }
295}
296
297/// Error returned when parsing a currency code.
298#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
299#[error("currency must contain exactly three uppercase ASCII letters")]
300pub struct CurrencyCodeParseError;
301
302impl FromStr for CurrencyCode {
303    type Err = CurrencyCodeParseError;
304
305    fn from_str(value: &str) -> Result<Self, Self::Err> {
306        if value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_uppercase()) {
307            Ok(Self(value.to_owned()))
308        } else {
309            Err(CurrencyCodeParseError)
310        }
311    }
312}
313
314impl fmt::Display for CurrencyCode {
315    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
316        formatter.write_str(&self.0)
317    }
318}
319
320impl Serialize for CurrencyCode {
321    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
322    where
323        S: Serializer,
324    {
325        serializer.serialize_str(&self.0)
326    }
327}
328
329impl<'de> Deserialize<'de> for CurrencyCode {
330    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
331    where
332        D: Deserializer<'de>,
333    {
334        String::deserialize(deserializer)?
335            .parse()
336            .map_err(serde::de::Error::custom)
337    }
338}
339
340/// Unit used by exact currency amounts.
341#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
342#[serde(rename_all = "snake_case")]
343pub enum CostUnit {
344    /// Amount is denominated in the currency's major unit, such as dollars.
345    MajorCurrency,
346}
347
348/// An exact persisted monetary amount.
349#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
350#[serde(rename_all = "camelCase")]
351pub struct ExactCost {
352    amount: DecimalAmount,
353    currency: CurrencyCode,
354    unit: CostUnit,
355}
356
357impl ExactCost {
358    /// Creates a cost denominated in a currency's major unit.
359    #[must_use]
360    pub const fn new(amount: DecimalAmount, currency: CurrencyCode) -> Self {
361        Self {
362            amount,
363            currency,
364            unit: CostUnit::MajorCurrency,
365        }
366    }
367
368    /// Returns the exact decimal amount.
369    #[must_use]
370    pub const fn amount(&self) -> &DecimalAmount {
371        &self.amount
372    }
373
374    /// Returns the currency code.
375    #[must_use]
376    pub const fn currency(&self) -> &CurrencyCode {
377        &self.currency
378    }
379
380    /// Returns the amount unit.
381    #[must_use]
382    pub const fn unit(&self) -> CostUnit {
383        self.unit
384    }
385}
tea_protocol/usage.rs

tea_protocol/
usage.rs
