llm_browser_testkit/
costs.rs

1//! Cost calculation, usage tracking, and pricing logic.
2
3use std::collections::{BTreeSet, HashMap};
4use std::sync::Mutex;
5
6use crate::endpoints::ResolvedEndpoint;
7
8/// Accumulated usage for a single endpoint.
9#[derive(Debug, Default, Clone)]
10pub struct EndpointUsage {
11    /// Number of calls made.
12    pub calls: u64,
13    /// Total input tokens consumed.
14    pub input_tokens: u64,
15    /// Total output tokens consumed.
16    pub output_tokens: u64,
17    /// Input tokens served from the provider's prompt cache.
18    pub cached_input_tokens: u64,
19    /// Input tokens written to the provider's prompt cache (cache creation).
20    pub cache_creation_input_tokens: u64,
21    /// Accumulated cost in USD.
22    pub cost: f64,
23    /// Model names observed on this endpoint, sorted and deduplicated.
24    pub models: BTreeSet<String>,
25}
26
27impl EndpointUsage {
28    const fn tokens(&self) -> u64 {
29        self.input_tokens + self.output_tokens
30    }
31}
32
33/// Aggregated usage across all endpoints for a test or scenario run.
34#[derive(Debug, Default, Clone)]
35pub struct UsageSnapshot {
36    /// Per-endpoint usage.
37    pub endpoints: HashMap<String, EndpointUsage>,
38    /// Total cost across all endpoints.
39    pub total_cost: f64,
40    /// Total calls across all endpoints.
41    pub total_calls: u64,
42    /// Total tokens across all endpoints.
43    pub total_tokens: u64,
44    /// Total input (prompt) tokens across all endpoints.
45    pub total_input_tokens: u64,
46    /// Total output (completion) tokens across all endpoints.
47    pub total_output_tokens: u64,
48    /// Total input tokens served from provider prompt caches.
49    pub total_cached_input_tokens: u64,
50    /// Total input tokens written to provider prompt caches.
51    pub total_cache_creation_input_tokens: u64,
52    /// Model names observed across all endpoints, sorted and deduplicated.
53    pub models: Vec<String>,
54}
55
56impl UsageSnapshot {
57    /// Creates a snapshot from per-endpoint usage data.
58    #[must_use]
59    pub fn from_endpoints(endpoints: &HashMap<String, EndpointUsage>) -> Self {
60        let total_cost = endpoints.values().map(|u| u.cost).sum();
61        let total_calls = endpoints.values().map(|u| u.calls).sum();
62        let total_tokens = endpoints.values().map(EndpointUsage::tokens).sum();
63        let total_input_tokens = endpoints.values().map(|u| u.input_tokens).sum();
64        let total_output_tokens = endpoints.values().map(|u| u.output_tokens).sum();
65        let total_cached_input_tokens = endpoints.values().map(|u| u.cached_input_tokens).sum();
66        let total_cache_creation_input_tokens = endpoints
67            .values()
68            .map(|u| u.cache_creation_input_tokens)
69            .sum();
70        let models: Vec<String> = endpoints
71            .values()
72            .flat_map(|u| u.models.iter().cloned())
73            .collect::<BTreeSet<_>>()
74            .into_iter()
75            .collect();
76        Self {
77            endpoints: endpoints.clone(),
78            total_cost,
79            total_calls,
80            total_tokens,
81            total_input_tokens,
82            total_output_tokens,
83            total_cached_input_tokens,
84            total_cache_creation_input_tokens,
85            models,
86        }
87    }
88}
89
90/// Thread-safe usage tracker for the test runner.
91pub struct UsageTracker {
92    inner: Mutex<UsageInner>,
93}
94
95struct UsageInner {
96    /// Per-endpoint usage for the current test.
97    per_endpoint: HashMap<String, EndpointUsage>,
98    /// Aggregated usage across all completed tests.
99    global: UsageSnapshot,
100    /// Per-test snapshots keyed by test name.
101    per_test: Vec<(String, UsageSnapshot)>,
102}
103
104impl UsageTracker {
105    /// Creates a new empty usage tracker.
106    #[must_use]
107    pub fn new() -> Self {
108        Self {
109            inner: Mutex::new(UsageInner {
110                per_endpoint: HashMap::new(),
111                global: UsageSnapshot::default(),
112                per_test: Vec::new(),
113            }),
114        }
115    }
116
117    /// Records a completed LLM call, adding usage, cost, and the answering
118    /// model to the endpoint's accumulator.
119    ///
120    /// # Panics
121    ///
122    /// Panics if the mutex is poisoned.
123    #[allow(clippy::significant_drop_tightening)]
124    pub fn record_llm_call(
125        &self,
126        endpoint_name: &str,
127        endpoint: &ResolvedEndpoint,
128        model: &str,
129        usage: &LlmUsage,
130    ) {
131        let cost = calculate_llm_cost(endpoint, usage.prompt_tokens, usage.completion_tokens);
132        let mut inner = self.inner.lock().unwrap();
133        let eu = inner
134            .per_endpoint
135            .entry(endpoint_name.to_owned())
136            .or_default();
137        eu.calls += 1;
138        eu.input_tokens += usage.prompt_tokens;
139        eu.output_tokens += usage.completion_tokens;
140        eu.cached_input_tokens += usage.cached_input_tokens;
141        eu.cache_creation_input_tokens += usage.cache_creation_input_tokens;
142        eu.cost += cost;
143        if !model.is_empty() {
144            eu.models.insert(model.to_owned());
145        }
146    }
147
148    /// Records a flat-cost call (MCP tool, agent task).
149    ///
150    /// # Panics
151    ///
152    /// Panics if the mutex is poisoned.
153    #[allow(clippy::significant_drop_tightening)]
154    pub fn record_flat_call(&self, endpoint_name: &str, endpoint: &ResolvedEndpoint) {
155        let mut inner = self.inner.lock().unwrap();
156        let eu = inner
157            .per_endpoint
158            .entry(endpoint_name.to_owned())
159            .or_default();
160        eu.calls += 1;
161        eu.cost += endpoint.per_call_price;
162    }
163
164    /// Reads current usage without locking for the full snapshot.
165    ///
166    /// # Panics
167    ///
168    /// Panics if the mutex is poisoned.
169    #[must_use]
170    pub fn current_test_snapshot(&self) -> UsageSnapshot {
171        let inner = self.inner.lock().unwrap();
172        UsageSnapshot::from_endpoints(&inner.per_endpoint)
173    }
174
175    /// Reads the global aggregated snapshot.
176    ///
177    /// # Panics
178    ///
179    /// Panics if the mutex is poisoned.
180    #[must_use]
181    pub fn global_snapshot(&self) -> UsageSnapshot {
182        let inner = self.inner.lock().unwrap();
183        inner.global.clone()
184    }
185
186    /// Reads per-test snapshots.
187    ///
188    /// # Panics
189    ///
190    /// Panics if the mutex is poisoned.
191    #[must_use]
192    pub fn per_test_snapshots(&self) -> Vec<(String, UsageSnapshot)> {
193        let inner = self.inner.lock().unwrap();
194        inner.per_test.clone()
195    }
196
197    /// Resets the per-test accumulator. Call at the start of each test.
198    ///
199    /// # Panics
200    ///
201    /// Panics if the mutex is poisoned.
202    pub fn reset_per_test(&self) {
203        let mut inner = self.inner.lock().unwrap();
204        inner.per_endpoint.clear();
205    }
206
207    /// Commits the current test's usage to the global accumulator and stores
208    /// it as a per-test snapshot.
209    ///
210    /// # Panics
211    ///
212    /// Panics if the mutex is poisoned.
213    pub fn commit_test(&self, test_name: &str) {
214        let mut inner = self.inner.lock().unwrap();
215        let snapshot = UsageSnapshot::from_endpoints(&inner.per_endpoint);
216        // Merge into global
217        let ep_snapshot = inner.per_endpoint.clone();
218        for (ep_name, ep_usage) in &ep_snapshot {
219            let ge = inner.global.endpoints.entry(ep_name.clone()).or_default();
220            ge.calls += ep_usage.calls;
221            ge.input_tokens += ep_usage.input_tokens;
222            ge.output_tokens += ep_usage.output_tokens;
223            ge.cached_input_tokens += ep_usage.cached_input_tokens;
224            ge.cache_creation_input_tokens += ep_usage.cache_creation_input_tokens;
225            ge.cost += ep_usage.cost;
226            ge.models.extend(ep_usage.models.iter().cloned());
227        }
228        inner.global.total_cost += snapshot.total_cost;
229        inner.global.total_calls += snapshot.total_calls;
230        inner.global.total_tokens += snapshot.total_tokens;
231        inner.global.total_input_tokens += snapshot.total_input_tokens;
232        inner.global.total_output_tokens += snapshot.total_output_tokens;
233        inner.global.total_cached_input_tokens += snapshot.total_cached_input_tokens;
234        inner.global.total_cache_creation_input_tokens +=
235            snapshot.total_cache_creation_input_tokens;
236        inner.global.models = inner
237            .global
238            .endpoints
239            .values()
240            .flat_map(|u| u.models.iter().cloned())
241            .collect::<BTreeSet<_>>()
242            .into_iter()
243            .collect();
244        inner.per_test.push((test_name.to_owned(), snapshot));
245    }
246}
247
248impl Default for UsageTracker {
249    fn default() -> Self {
250        Self::new()
251    }
252}
253
254/// Calculates the cost of an LLM call based on token pricing.
255#[allow(clippy::cast_precision_loss, clippy::suboptimal_flops)]
256#[must_use]
257pub fn calculate_llm_cost(
258    endpoint: &ResolvedEndpoint,
259    input_tokens: u64,
260    output_tokens: u64,
261) -> f64 {
262    let input_cost = (input_tokens as f64 / 1_000_000.0) * endpoint.input_price_per_1m;
263    let output_cost = (output_tokens as f64 / 1_000_000.0) * endpoint.output_price_per_1m;
264    input_cost + output_cost
265}
266
267/// Usage info extracted from an LLM API response.
268#[derive(Debug, Default, Clone, Copy)]
269pub struct LlmUsage {
270    /// Number of prompt / input tokens.
271    pub prompt_tokens: u64,
272    /// Number of completion / output tokens.
273    pub completion_tokens: u64,
274    /// Total tokens used.
275    pub total_tokens: u64,
276    /// Input tokens served from the provider's prompt cache (cache hit).
277    pub cached_input_tokens: u64,
278    /// Input tokens written to the provider's prompt cache (cache creation).
279    pub cache_creation_input_tokens: u64,
280}
281
282/// Result of an LLM chat call including usage data.
283#[derive(Debug, Clone)]
284pub struct LlmResponse {
285    /// The message content from the LLM.
286    pub content: String,
287    /// Token usage from the API response.
288    pub usage: LlmUsage,
289}
290
291/// Extracts token usage from an OpenAI-compatible API response JSON.
292///
293/// Cached prompt tokens are read from the `OpenAI`/`OpenRouter`
294/// `usage.prompt_tokens_details.cached_tokens` field, falling back to the
295/// `Anthropic`-compatible `usage.cache_read_input_tokens` and `DeepSeek`
296/// `usage.prompt_cache_hit_tokens` spellings. Cache-creation (write) tokens
297/// are read from the `Anthropic`-compatible
298/// `usage.cache_creation_input_tokens`. Providers that do not report prompt
299/// caching yield `0`.
300///
301/// Note the provider semantics differ: for `OpenAI`-compatible responses
302/// `cached_tokens` is a subset of `prompt_tokens`, whereas for
303/// `Anthropic`/Bedrock-style responses the cache counters are reported in
304/// addition to `inputTokens`.
305#[must_use]
306pub fn extract_usage(value: &serde_json::Value) -> LlmUsage {
307    let usage = &value["usage"];
308    LlmUsage {
309        prompt_tokens: usage["prompt_tokens"].as_u64().unwrap_or(0),
310        completion_tokens: usage["completion_tokens"].as_u64().unwrap_or(0),
311        total_tokens: usage["total_tokens"].as_u64().unwrap_or(0),
312        cached_input_tokens: usage["prompt_tokens_details"]["cached_tokens"]
313            .as_u64()
314            .or_else(|| usage["cache_read_input_tokens"].as_u64())
315            .or_else(|| usage["prompt_cache_hit_tokens"].as_u64())
316            .unwrap_or(0),
317        cache_creation_input_tokens: usage["cache_creation_input_tokens"].as_u64().unwrap_or(0),
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use crate::costs::{calculate_llm_cost, LlmUsage, UsageTracker};
324    use crate::endpoints::ResolvedEndpoint;
325    use crate::scenario::EndpointType;
326
327    fn make_endpoint(
328        name: &str,
329        input_price: f64,
330        output_price: f64,
331        per_call: f64,
332    ) -> ResolvedEndpoint {
333        ResolvedEndpoint {
334            name: name.to_owned(),
335            endpoint_type: EndpointType::Llm,
336            url: String::new(),
337            model: None,
338            api_key: None,
339            headers: std::collections::HashMap::new(),
340            command: None,
341            args: vec![],
342            vision: false,
343            input_price_per_1m: input_price,
344            output_price_per_1m: output_price,
345            per_call_price: per_call,
346            max_attempts: 3,
347            fallbacks: vec![],
348            provider: crate::scenario::Provider::Openai,
349            deployment: None,
350            api_version: None,
351            auth: crate::scenario::AuthConfig::default(),
352            header_commands: std::collections::HashMap::new(),
353            aws: crate::scenario::AwsConfig::default(),
354        }
355    }
356
357    fn usage(prompt: u64, completion: u64, cached: u64) -> LlmUsage {
358        LlmUsage {
359            prompt_tokens: prompt,
360            completion_tokens: completion,
361            total_tokens: prompt + completion,
362            cached_input_tokens: cached,
363            cache_creation_input_tokens: 0,
364        }
365    }
366
367    #[test]
368    fn test_calculate_llm_cost() {
369        let ep = make_endpoint("test", 0.15, 0.60, 0.0);
370        // 1M input tokens = $0.15, 500K output = $0.30
371        let cost = calculate_llm_cost(&ep, 1_000_000, 500_000);
372        assert!((cost - 0.45).abs() < 0.001);
373    }
374
375    #[test]
376    fn test_calculate_zero_cost() {
377        let ep = make_endpoint("free", 0.0, 0.0, 0.0);
378        let cost = calculate_llm_cost(&ep, 1_000_000, 1_000_000);
379        assert!((cost - 0.0).abs() < f64::EPSILON);
380    }
381
382    #[test]
383    fn test_usage_tracker_record_llm() {
384        let tracker = UsageTracker::new();
385        let ep = make_endpoint("gpt4", 2.50, 10.0, 0.0);
386        tracker.record_llm_call("gpt4", &ep, "gpt-4o", &usage(1000, 500, 200));
387
388        let snap = tracker.current_test_snapshot();
389        assert_eq!(snap.total_calls, 1);
390        assert_eq!(snap.total_tokens, 1500);
391        assert_eq!(snap.total_input_tokens, 1000);
392        assert_eq!(snap.total_output_tokens, 500);
393        assert_eq!(snap.total_cached_input_tokens, 200);
394        assert_eq!(snap.models, vec!["gpt-4o".to_owned()]);
395        assert!(
396            snap.total_cost > 0.0,
397            "expected cost > 0, got {}",
398            snap.total_cost
399        );
400
401        let ep_usage = snap.endpoints.get("gpt4").unwrap();
402        assert_eq!(ep_usage.calls, 1);
403        assert_eq!(ep_usage.input_tokens, 1000);
404        assert_eq!(ep_usage.output_tokens, 500);
405        assert_eq!(ep_usage.cached_input_tokens, 200);
406        assert!(ep_usage.models.contains("gpt-4o"));
407    }
408
409    #[test]
410    fn test_usage_tracker_record_flat() {
411        let tracker = UsageTracker::new();
412        let ep = make_endpoint("agent", 0.0, 0.0, 0.01);
413        tracker.record_flat_call("agent", &ep);
414        tracker.record_flat_call("agent", &ep);
415
416        let snap = tracker.current_test_snapshot();
417        assert_eq!(snap.total_calls, 2);
418        assert!((snap.total_cost - 0.02).abs() < f64::EPSILON);
419    }
420
421    #[test]
422    fn test_usage_tracker_multiple_endpoints() {
423        let tracker = UsageTracker::new();
424        let fast = make_endpoint("fast", 0.15, 0.60, 0.0);
425        let slow = make_endpoint("slow", 2.50, 10.0, 0.0);
426
427        tracker.record_llm_call("fast", &fast, "fast-model", &usage(100, 50, 0));
428        tracker.record_llm_call("slow", &slow, "slow-model", &usage(200, 100, 0));
429
430        let snap = tracker.current_test_snapshot();
431        assert_eq!(snap.total_calls, 2);
432        assert_eq!(snap.endpoints.len(), 2);
433        assert_eq!(
434            snap.models,
435            vec!["fast-model".to_owned(), "slow-model".to_owned()]
436        );
437    }
438
439    #[test]
440    fn test_usage_tracker_reset_and_commit() {
441        let tracker = UsageTracker::new();
442        let ep = make_endpoint("test", 0.15, 0.60, 0.0);
443
444        tracker.record_llm_call("test", &ep, "m1", &usage(100, 50, 0));
445        tracker.commit_test("test1");
446        tracker.reset_per_test();
447
448        tracker.record_llm_call("test", &ep, "m2", &usage(200, 100, 0));
449        tracker.commit_test("test2");
450
451        let global = tracker.global_snapshot();
452        assert_eq!(global.total_calls, 2);
453        assert_eq!(global.total_tokens, 450);
454        assert_eq!(global.models, vec!["m1".to_owned(), "m2".to_owned()]);
455
456        let per_test = tracker.per_test_snapshots();
457        assert_eq!(per_test.len(), 2);
458        assert_eq!(per_test[0].0, "test1");
459        assert_eq!(per_test[1].0, "test2");
460    }
461}
llm_browser_testkit/costs.rs

llm_browser_testkit/
costs.rs
