Skip to main content

kernel/model/
usage.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4// Copyright (c) 2026 2youg1 and the sprawling contributors
5
6//! What a provider reports a call cost in tokens, and how a stored row
7//! of it reads under each meaning `input_tokens` has had.
8//! The part `crates/kernel/spec/Model.lean` specifies this module with the
9//! rest of `kernel::model`.
10
11use serde::{Deserialize, Serialize};
12
13use crate::budget::Tokens;
14use crate::model::wire::DialectKind;
15
16/// One cache count as a provider reported it: a number, or nothing at all.
17///
18/// A provider that never reports its cache and one that reports a zero
19/// hit are different facts to the User, so the two do not share a value
20/// (`crates/kernel/spec/Model.lean` D36). A stored row spells
21/// `Unreported` by leaving the key out; an absent key or `null` reads
22/// back as `Unreported`, a present `0` as `Reported(0)`.
23#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
24pub enum CacheCount {
25    Reported(Tokens),
26    #[default]
27    Unreported,
28}
29
30impl CacheCount {
31    /// The count when the provider gave one.
32    #[must_use]
33    pub const fn reported(self) -> Option<Tokens> {
34        match self {
35            CacheCount::Reported(tokens) => Some(tokens),
36            CacheCount::Unreported => None,
37        }
38    }
39
40    /// The count for arithmetic that subtracts the cached part from the
41    /// whole prompt (pricing, the context gauge): an unreported count is
42    /// read as nothing cached, the reading those sums had before the
43    /// distinction existed.
44    #[must_use]
45    pub const fn or_zero(self) -> Tokens {
46        match self {
47            CacheCount::Reported(tokens) => tokens,
48            CacheCount::Unreported => Tokens::new(0),
49        }
50    }
51}
52
53impl From<Option<Tokens>> for CacheCount {
54    fn from(count: Option<Tokens>) -> CacheCount {
55        match count {
56            Some(tokens) => CacheCount::Reported(tokens),
57            None => CacheCount::Unreported,
58        }
59    }
60}
61
62/// Provider-reported token counts.
63///
64/// `input_tokens` is every prompt token of the request, cached or not,
65/// whichever dialect reported it; `cache_read_tokens` and
66/// `cache_write_tokens` are parts of it, each [`CacheCount::Unreported`]
67/// when the provider gave no number. `dialect` is `None` for a
68/// count no provider wire produced (scripted models, tests).
69///
70/// Serialised rows carry the meaning's version as `v`. A row without `v` was
71/// written before the one meaning existed and names no dialect, so
72/// deserialisation reads it as Anthropic's uncached count exactly when
73/// its cache parts exceed it, which no whole-prompt count can do.
74#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
75#[serde(into = "UsageRow", try_from = "UsageRow")]
76pub struct ModelUsage {
77    pub input_tokens: Tokens,
78    pub output_tokens: Tokens,
79    pub cache_read_tokens: CacheCount,
80    pub cache_write_tokens: CacheCount,
81    pub dialect: Option<DialectKind>,
82}
83
84/// The version of the meaning `ModelUsage.input_tokens` has in a stored
85/// row: the whole prompt. The ledger is append-only, so a row written
86/// under an older meaning keeps its bytes and is converted on read.
87const USAGE_MEANING: u8 = 1;
88
89/// `ModelUsage` as a row spells it: the version travels in the row
90/// rather than in the type, because only a reader of old bytes needs it.
91#[derive(Serialize, Deserialize)]
92struct UsageRow {
93    input_tokens: Tokens,
94    output_tokens: Tokens,
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    cache_read_tokens: Option<Tokens>,
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    cache_write_tokens: Option<Tokens>,
99    #[serde(default, skip_serializing_if = "Option::is_none")]
100    v: Option<u8>,
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    dialect: Option<DialectKind>,
103}
104
105impl From<ModelUsage> for UsageRow {
106    fn from(usage: ModelUsage) -> UsageRow {
107        UsageRow {
108            input_tokens: usage.input_tokens,
109            output_tokens: usage.output_tokens,
110            cache_read_tokens: usage.cache_read_tokens.reported(),
111            cache_write_tokens: usage.cache_write_tokens.reported(),
112            v: Some(USAGE_MEANING),
113            dialect: usage.dialect,
114        }
115    }
116}
117
118impl TryFrom<UsageRow> for ModelUsage {
119    type Error = String;
120
121    fn try_from(row: UsageRow) -> Result<ModelUsage, String> {
122        let input_tokens = match row.v {
123            Some(USAGE_MEANING) => row.input_tokens,
124            Some(other) => {
125                return Err(format!(
126                    "usage row version {other} is newer than this build reads \
127                     (it reads {USAGE_MEANING}); read the ledger with a newer sprawling"
128                ));
129            }
130            None => {
131                let cached = CacheCount::from(row.cache_read_tokens)
132                    .or_zero()
133                    .checked_add(CacheCount::from(row.cache_write_tokens).or_zero())
134                    .ok_or("usage row cache counts overflow a token count")?;
135                if cached > row.input_tokens {
136                    row.input_tokens
137                        .checked_add(cached)
138                        .ok_or("usage row counts overflow a token count")?
139                } else {
140                    row.input_tokens
141                }
142            }
143        };
144        Ok(ModelUsage {
145            input_tokens,
146            output_tokens: row.output_tokens,
147            cache_read_tokens: CacheCount::from(row.cache_read_tokens),
148            cache_write_tokens: CacheCount::from(row.cache_write_tokens),
149            dialect: row.dialect,
150        })
151    }
152}