storage/attribution/
report.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//! Where the money went. Four independent cuts of
7//! one authoritative total: by Run, by actor, by prefix segment, by
8//! tool.
9//!
10//! The invariant that makes the report trustworthy (A20) is that each
11//! cut sums to exactly the billed total — not approximately, not after
12//! rounding. Percentages cannot promise that, so nothing here is a
13//! percentage: shares are split by the largest-remainder method, which
14//! distributes the floor divisions and then hands the leftover units
15//! out one at a time, in a fixed order. What is billed is what is
16//! attributed, to the microdollar.
17//!
18//! The total comes from `model_returned.billed_usd_micros` and nowhere
19//! else. This module never prices a call — that authority is
20//! `gateway::cost`, and a second one would be a second answer.
21//!
22//! Specified by `crates/storage/spec/Attribution.lean` §8-7.
23
24//! Attribution reports: what each wave billed.
25
26use super::split::{add, quantify, segment_weights, split};
27
28/// The bucket a call falls into when the ledger gives no basis for a
29/// finer split — an honest "unsplit", never a silent drop.
30pub(crate) const NO_TOOL: &str = "no_tool";
31pub(crate) const NO_SKILL: &str = "no_skill";
32pub(crate) const NO_SEGMENT: &str = "unattributed";
33pub(crate) const WINDOW_SLOT: &str = "window";
34
35use std::collections::BTreeMap;
36
37use kernel::{EventKind, EventRecord, RunId, UsdMicros};
38use serde_json::Value;
39
40use crate::error::StorageError;
41
42#[derive(Clone, Default, serde::Serialize, serde::Deserialize)]
43pub struct Attribution {
44    by_run: BTreeMap<String, u64>,
45    by_actor: BTreeMap<String, u64>,
46    by_segment: BTreeMap<String, u64>,
47    by_tool: BTreeMap<String, u64>,
48    by_skill: BTreeMap<String, u64>,
49    total: u64,
50    /// Basis for the next model_returned: each run's most recent prefix
51    /// shape, and the tool bytes returned since the previous call. The
52    /// shape is keyed by run because a run records its prompt once, so
53    /// the latest `prompt_assembled` in the ledger may be another run's.
54    segment_weights: BTreeMap<RunId, Vec<(String, u64)>>,
55    tool_weights: BTreeMap<String, u64>,
56    skill_weights: BTreeMap<String, u64>,
57    unpriced: Unpriced,
58}
59
60pub struct AttributionReport {
61    pub total: UsdMicros,
62    pub by_run: Vec<(String, UsdMicros)>,
63    pub by_actor: Vec<(String, UsdMicros)>,
64    pub by_segment: Vec<(String, UsdMicros)>,
65    pub by_tool: Vec<(String, UsdMicros)>,
66    pub by_skill: Vec<(String, UsdMicros)>,
67    pub unpriced: Unpriced,
68}
69
70/// The model calls that came back with no authoritative amount. They
71/// add nothing to `total`, so without this count a city whose provider
72/// never prices a call reads the same as a city that never ran one; the
73/// tokens are the one measure of their use the ledger does hold.
74#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
75pub struct Unpriced {
76    pub calls: u64,
77    pub tokens: u64,
78}
79
80impl Attribution {
81    pub fn new() -> Attribution {
82        Attribution::default()
83    }
84
85    /// Folds one record in. Only `model_returned` moves money; two other
86    /// kinds set the basis on which the next call is split, and
87    /// `run_frozen` drops the basis of a run that makes no further call.
88    #[expect(
89        clippy::wildcard_enum_match_arm,
90        reason = "four kinds carry money or its basis; the rest of the vocabulary is silent here"
91    )]
92    pub fn apply(&mut self, record: &EventRecord) -> Result<(), StorageError> {
93        match record.kind() {
94            EventKind::PromptAssembled => {
95                self.segment_weights
96                    .insert(record.run(), segment_weights(record.data().as_map()));
97            }
98            EventKind::RunFrozen => {
99                self.segment_weights.remove(&record.run());
100            }
101            EventKind::ToolResult => {
102                let name = record
103                    .data()
104                    .as_map()
105                    .get("name")
106                    .and_then(Value::as_str)
107                    .unwrap_or("unnamed")
108                    .to_owned();
109                let bytes = record
110                    .data()
111                    .as_map()
112                    .get("bytes")
113                    .and_then(Value::as_u64)
114                    .unwrap_or(1);
115                let slot = self.tool_weights.entry(name).or_insert(0);
116                *slot = slot.saturating_add(bytes);
117                // The SKILL a call was made under, when one was; a call
118                // without one goes to the `no_skill` bucket.
119                let skill = record
120                    .data()
121                    .as_map()
122                    .get("skill")
123                    .and_then(Value::as_str)
124                    .unwrap_or(NO_SKILL)
125                    .to_owned();
126                let slot = self.skill_weights.entry(skill).or_insert(0);
127                *slot = slot.saturating_add(bytes);
128            }
129            EventKind::ModelReturned => {
130                let Some(billed) = record
131                    .data()
132                    .as_map()
133                    .get("billed_usd_micros")
134                    .and_then(Value::as_u64)
135                else {
136                    // A call with no authoritative amount attributes
137                    // nothing. Estimating here would invent money.
138                    self.unpriced.calls = self.unpriced.calls.saturating_add(1);
139                    self.unpriced.tokens = self
140                        .unpriced
141                        .tokens
142                        .saturating_add(usage_tokens(record.data().as_map()));
143                    self.tool_weights.clear();
144                    self.skill_weights.clear();
145                    return Ok(());
146                };
147                self.total = self.total.saturating_add(billed);
148                add(&mut self.by_run, &record.run().to_string(), billed);
149                add(&mut self.by_actor, record.who(), billed);
150                let segment_basis = self
151                    .segment_weights
152                    .get(&record.run())
153                    .map_or(&[][..], Vec::as_slice);
154                for (bucket, share) in split(billed, segment_basis, NO_SEGMENT) {
155                    add(&mut self.by_segment, &bucket, share);
156                }
157                let tool_basis: Vec<(String, u64)> = self
158                    .tool_weights
159                    .iter()
160                    .map(|(k, v)| (k.clone(), *v))
161                    .collect();
162                for (bucket, share) in split(billed, &tool_basis, NO_TOOL) {
163                    add(&mut self.by_tool, &bucket, share);
164                }
165                let skill_basis: Vec<(String, u64)> = self
166                    .skill_weights
167                    .iter()
168                    .map(|(k, v)| (k.clone(), *v))
169                    .collect();
170                for (bucket, share) in split(billed, &skill_basis, NO_SKILL) {
171                    add(&mut self.by_skill, &bucket, share);
172                }
173                // The wave is settled; the next call starts a new basis.
174                self.tool_weights.clear();
175                self.skill_weights.clear();
176            }
177            _ => {}
178        }
179        Ok(())
180    }
181
182    /// What `run` was billed, when this fold holds a row for it: one
183    /// lookup rather than a whole `report()` per question.
184    pub fn billed_to(&self, run: &RunId) -> Option<UsdMicros> {
185        self.by_run
186            .get(&run.to_string())
187            .copied()
188            .map(UsdMicros::new)
189    }
190
191    /// Drops the `by_run` row of every run `keep` refuses. `total` and
192    /// the other four cuts keep the money, so they still sum to `total`
193    /// while `by_run` shrinks to the runs a caller still holds warm.
194    pub fn retain_runs(&mut self, keep: impl Fn(&RunId) -> bool) {
195        self.by_run
196            .retain(|name, _| RunId::parse(name).map_or(true, |run| keep(&run)));
197    }
198
199    pub fn report(&self) -> AttributionReport {
200        AttributionReport {
201            total: UsdMicros::new(self.total),
202            by_run: quantify(&self.by_run),
203            by_actor: quantify(&self.by_actor),
204            by_segment: quantify(&self.by_segment),
205            by_tool: quantify(&self.by_tool),
206            by_skill: quantify(&self.by_skill),
207            unpriced: self.unpriced,
208        }
209    }
210}
211
212/// The four token counts of one `model_returned`'s `usage`, summed; a
213/// call that reported no usage used no tokens anyone measured.
214fn usage_tokens(data: &serde_json::Map<String, Value>) -> u64 {
215    let usage = data.get("usage");
216    [
217        "input_tokens",
218        "output_tokens",
219        "cache_read_tokens",
220        "cache_write_tokens",
221    ]
222    .into_iter()
223    .filter_map(|field| usage.and_then(|u| u.get(field)).and_then(Value::as_u64))
224    .fold(0, u64::saturating_add)
225}
226
227#[cfg(test)]
228mod tests;
storage/attribution/report.rs

storage/attribution/
report.rs
