leviath_runtime/
inference_usage.rs

1//! What one provider call cost, recorded where it lands.
2//!
3//! A run bills for four kinds of call and, until this module existed, counted
4//! one of them. Stage turns went into [`TokenTotals`];
5//! the compaction, title, and routing lanes each threw their usage away at the
6//! outcome boundary, because each channel carried only the payload its collector
7//! wanted - a summary, a title, a stage name - and usage was not it.
8//!
9//! Two things follow from putting the accounting in one place rather than
10//! repeating it per lane. The cumulative totals finally cover every call, so a
11//! run's reported spend is what the provider actually billed. And each call
12//! writes a [`RunRecord::InferenceUsage`] as it lands, which is the part
13//! [`RunRecord::Progress`] cannot do: progress counters are cumulative, so two
14//! calls between two ticks arrive as their sum, and a chart of that sum shows a
15//! spike no single call ever made.
16
17use leviath_core::run_archive::{InferenceKind, RunRecord};
18
19use crate::persistence::{RunMetadata, TokenTotals};
20use crate::persistence_bridge::PersistMsg;
21use crate::pipeline::PersistenceStage;
22
23/// Everything one call needs to report itself.
24///
25/// Passed as a struct rather than eight arguments because the four token counts
26/// are trivially transposable at a call site and the compiler would not catch
27/// it.
28pub struct CallUsage<'a> {
29    /// Which kind of call this was.
30    pub kind: InferenceKind,
31    /// The stage the run was in; empty for the title call, which has none.
32    pub stage: &'a str,
33    /// The stage-local iteration index.
34    pub iteration: usize,
35    /// The provider that served the call.
36    pub provider: &'a str,
37    /// The model the call targeted.
38    pub model: &'a str,
39    /// What the provider billed.
40    pub usage: &'a leviath_providers::TokenUsage,
41}
42
43/// Fold one call into the run's cumulative totals and journal what it cost.
44///
45/// Both halves are optional and independent: a world with no `TokenTotals` (a
46/// bare test agent) still journals, and a world with no persistence lane or run
47/// metadata - tests, unpersisted agents - still counts. Neither absence is an
48/// error, which is why this takes options rather than making callers branch.
49pub fn record_call(
50    totals: Option<&mut TokenTotals>,
51    persist: Option<&PersistenceStage>,
52    metadata: Option<&RunMetadata>,
53    call: &CallUsage<'_>,
54) {
55    if let Some(totals) = totals {
56        totals.add_usage(call.usage);
57    }
58    let (Some(persist), Some(md)) = (persist, metadata) else {
59        return;
60    };
61    let record = RunRecord::InferenceUsage {
62        kind: call.kind,
63        stage: call.stage.to_string(),
64        iteration: call.iteration,
65        provider: call.provider.to_string(),
66        model: call.model.to_string(),
67        prompt_tokens: call.usage.prompt_tokens,
68        completion_tokens: call.usage.completion_tokens,
69        cached_tokens: call.usage.cached_tokens,
70        cache_write_tokens: call.usage.cache_write_tokens,
71        at: chrono::Utc::now().timestamp(),
72    };
73    // No ack: a usage record is telemetry, and nothing downstream waits on it
74    // the way the tool lane waits on its batch record being durable before
75    // anything can run.
76    let _ = persist.0.send(PersistMsg::Append {
77        run_id: md.run_id.clone(),
78        record: Box::new(record),
79        ack: None,
80    });
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    fn usage() -> leviath_providers::TokenUsage {
88        leviath_providers::TokenUsage {
89            prompt_tokens: 100,
90            completion_tokens: 20,
91            cached_tokens: 3,
92            cache_write_tokens: 4,
93            total_tokens: 120,
94        }
95    }
96
97    fn metadata() -> RunMetadata {
98        RunMetadata {
99            run_id: "run-u".to_string(),
100            agent_name: "a".to_string(),
101            agent_path: "/a".to_string(),
102            task: "t".to_string(),
103            model: None,
104            workdir: "/w".to_string(),
105            num_stages: 1,
106            started_at: 0,
107            parent_run_id: None,
108            metadata: Default::default(),
109            callback_url: None,
110            callback_secret: None,
111            title: None,
112            unattended: false,
113            read_paths: None,
114            output_request: None,
115        }
116    }
117
118    fn call(kind: InferenceKind, u: &leviath_providers::TokenUsage) -> CallUsage<'_> {
119        CallUsage {
120            kind,
121            stage: "plan",
122            iteration: 2,
123            provider: "anthropic",
124            model: "claude-sonnet-5",
125            usage: u,
126        }
127    }
128
129    /// The two halves are independent by design, so the four combinations of
130    /// "has totals" and "has a journal" all have to behave. A world missing
131    /// either is ordinary - tests and unpersisted agents run that way - and an
132    /// absence must not cost the half that is present.
133    #[test]
134    fn counting_and_journaling_are_independent() {
135        let u = usage();
136
137        // Both present: counted and written.
138        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
139        let tx_for_noise = tx.clone();
140        let mut totals = TokenTotals::default();
141        record_call(
142            Some(&mut totals),
143            Some(&PersistenceStage(tx)),
144            Some(&metadata()),
145            &call(InferenceKind::Compaction, &u),
146        );
147        assert_eq!(totals.prompt_tokens, 100);
148        // The persistence channel carries snapshots and buffered log lines on
149        // the same wire, so the drain has to pick ours out of mixed traffic
150        // rather than assume the next message is it.
151        let _ = tx_for_noise.send(crate::persistence_bridge::PersistMsg::StageLines {
152            run_id: "run-u".to_string(),
153            output_appends: vec![],
154            log_appends: vec![],
155        });
156        let mut appended: Vec<(String, RunRecord)> = Vec::new();
157        while let Ok(msg) = rx.try_recv() {
158            if let crate::persistence_bridge::PersistMsg::Append { run_id, record, .. } = msg {
159                appended.push((run_id, *record));
160            }
161        }
162        assert_eq!(appended.len(), 1, "one call, one record");
163        let (run_id, record) = appended.remove(0);
164        assert_eq!(run_id, "run-u");
165        // Asserted on the serialized form with the wall-clock stamp lifted out,
166        // so this pins the field names a journal reader parses without pinning
167        // the one value that cannot be known ahead of time.
168        let mut value = serde_json::to_value(&record).unwrap();
169        let fields = value["InferenceUsage"].as_object_mut().unwrap();
170        assert!(fields.remove("at").is_some(), "a call is stamped");
171        assert_eq!(
172            value,
173            serde_json::json!({
174                "InferenceUsage": {
175                    "kind": "compaction",
176                    "stage": "plan",
177                    "iteration": 2,
178                    "provider": "anthropic",
179                    "model": "claude-sonnet-5",
180                    "prompt_tokens": 100,
181                    "completion_tokens": 20,
182                    "cached_tokens": 3,
183                    "cache_write_tokens": 4,
184                }
185            })
186        );
187
188        // No journal: still counted.
189        let mut totals = TokenTotals::default();
190        record_call(
191            Some(&mut totals),
192            None,
193            Some(&metadata()),
194            &call(InferenceKind::Stage, &u),
195        );
196        assert_eq!(totals.prompt_tokens, 100);
197
198        // A journal but no run metadata: nothing to address the record to, so
199        // nothing is written - and the count still lands.
200        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
201        let mut totals = TokenTotals::default();
202        record_call(
203            Some(&mut totals),
204            Some(&PersistenceStage(tx)),
205            None,
206            &call(InferenceKind::Title, &u),
207        );
208        assert_eq!(totals.prompt_tokens, 100);
209        assert!(rx.try_recv().is_err());
210
211        // Neither: a no-op that must not panic.
212        record_call(None, None, None, &call(InferenceKind::Routing, &u));
213    }
214
215    /// Totals accumulate across calls rather than being overwritten - the bug
216    /// that made a run report its last call instead of its bill would pass a
217    /// single-call test.
218    #[test]
219    fn repeated_calls_accumulate() {
220        let u = usage();
221        let mut totals = TokenTotals::default();
222        for _ in 0..3 {
223            record_call(
224                Some(&mut totals),
225                None,
226                None,
227                &call(InferenceKind::Stage, &u),
228            );
229        }
230        assert_eq!(totals.prompt_tokens, 300);
231        assert_eq!(totals.completion_tokens, 60);
232        assert_eq!(totals.cached_tokens, 9);
233        assert_eq!(totals.cache_write_tokens, 12);
234    }
235}
leviath_runtime/inference_usage.rs

leviath_runtime/
inference_usage.rs
