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/// Provider-reported token counts; absent wire fields read as zero.
17///
18/// `input_tokens` is every prompt token of the request, cached or not,
19/// whichever dialect reported it; `cache_read_tokens` and
20/// `cache_write_tokens` are parts of it. `dialect` is `None` for a
21/// count no provider wire produced (scripted models, tests).
22///
23/// Serialised rows carry the meaning's version as `v`. A row without `v` was
24/// written before the one meaning existed and names no dialect, so
25/// deserialisation reads it as Anthropic's uncached count exactly when
26/// its cache parts exceed it, which no whole-prompt count can do.
27#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(into = "UsageRow", try_from = "UsageRow")]
29pub struct ModelUsage {
30    pub input_tokens: Tokens,
31    pub output_tokens: Tokens,
32    pub cache_read_tokens: Tokens,
33    pub cache_write_tokens: Tokens,
34    pub dialect: Option<DialectKind>,
35}
36
37/// The version of the meaning `ModelUsage.input_tokens` has in a stored
38/// row: the whole prompt. The ledger is append-only, so a row written
39/// under an older meaning keeps its bytes and is converted on read.
40const USAGE_MEANING: u8 = 1;
41
42/// `ModelUsage` as a row spells it: the version travels in the row
43/// rather than in the type, because only a reader of old bytes needs it.
44#[derive(Serialize, Deserialize)]
45struct UsageRow {
46    input_tokens: Tokens,
47    output_tokens: Tokens,
48    cache_read_tokens: Tokens,
49    cache_write_tokens: Tokens,
50    #[serde(default, skip_serializing_if = "Option::is_none")]
51    v: Option<u8>,
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    dialect: Option<DialectKind>,
54}
55
56impl From<ModelUsage> for UsageRow {
57    fn from(usage: ModelUsage) -> UsageRow {
58        UsageRow {
59            input_tokens: usage.input_tokens,
60            output_tokens: usage.output_tokens,
61            cache_read_tokens: usage.cache_read_tokens,
62            cache_write_tokens: usage.cache_write_tokens,
63            v: Some(USAGE_MEANING),
64            dialect: usage.dialect,
65        }
66    }
67}
68
69impl TryFrom<UsageRow> for ModelUsage {
70    type Error = String;
71
72    fn try_from(row: UsageRow) -> Result<ModelUsage, String> {
73        let input_tokens = match row.v {
74            Some(USAGE_MEANING) => row.input_tokens,
75            Some(other) => {
76                return Err(format!(
77                    "usage row version {other} is newer than this build reads \
78                     (it reads {USAGE_MEANING}); read the ledger with a newer sprawling"
79                ));
80            }
81            None => {
82                let cached = row
83                    .cache_read_tokens
84                    .checked_add(row.cache_write_tokens)
85                    .ok_or("usage row cache counts overflow a token count")?;
86                if cached > row.input_tokens {
87                    row.input_tokens
88                        .checked_add(cached)
89                        .ok_or("usage row counts overflow a token count")?
90                } else {
91                    row.input_tokens
92                }
93            }
94        };
95        Ok(ModelUsage {
96            input_tokens,
97            output_tokens: row.output_tokens,
98            cache_read_tokens: row.cache_read_tokens,
99            cache_write_tokens: row.cache_write_tokens,
100            dialect: row.dialect,
101        })
102    }
103}