llm_browser_testkit/
costs.rs

1//! Cost calculation, usage tracking, and pricing logic.
2
3use std::collections::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    /// Accumulated cost in USD.
18    pub cost: f64,
19}
20
21impl EndpointUsage {
22    const fn tokens(&self) -> u64 {
23        self.input_tokens + self.output_tokens
24    }
25}
26
27/// Aggregated usage across all endpoints for a test or scenario run.
28#[derive(Debug, Default, Clone)]
29pub struct UsageSnapshot {
30    /// Per-endpoint usage.
31    pub endpoints: HashMap<String, EndpointUsage>,
32    /// Total cost across all endpoints.
33    pub total_cost: f64,
34    /// Total calls across all endpoints.
35    pub total_calls: u64,
36    /// Total tokens across all endpoints.
37    pub total_tokens: u64,
38}
39
40impl UsageSnapshot {
41    /// Creates a snapshot from per-endpoint usage data.
42    #[must_use]
43    pub fn from_endpoints(endpoints: &HashMap<String, EndpointUsage>) -> Self {
44        let total_cost = endpoints.values().map(|u| u.cost).sum();
45        let total_calls = endpoints.values().map(|u| u.calls).sum();
46        let total_tokens = endpoints.values().map(EndpointUsage::tokens).sum();
47        Self {
48            endpoints: endpoints.clone(),
49            total_cost,
50            total_calls,
51            total_tokens,
52        }
53    }
54}
55
56/// Thread-safe usage tracker for the test runner.
57pub struct UsageTracker {
58    inner: Mutex<UsageInner>,
59}
60
61struct UsageInner {
62    /// Per-endpoint usage for the current test.
63    per_endpoint: HashMap<String, EndpointUsage>,
64    /// Aggregated usage across all completed tests.
65    global: UsageSnapshot,
66    /// Per-test snapshots keyed by test name.
67    per_test: Vec<(String, UsageSnapshot)>,
68}
69
70impl UsageTracker {
71    /// Creates a new empty usage tracker.
72    #[must_use]
73    pub fn new() -> Self {
74        Self {
75            inner: Mutex::new(UsageInner {
76                per_endpoint: HashMap::new(),
77                global: UsageSnapshot::default(),
78                per_test: Vec::new(),
79            }),
80        }
81    }
82
83    /// Records a completed call, adding usage and cost.
84    ///
85    /// # Panics
86    ///
87    /// Panics if the mutex is poisoned.
88    #[allow(clippy::significant_drop_tightening)]
89    pub fn record_llm_call(
90        &self,
91        endpoint_name: &str,
92        endpoint: &ResolvedEndpoint,
93        input_tokens: u64,
94        output_tokens: u64,
95    ) {
96        let cost = calculate_llm_cost(endpoint, input_tokens, output_tokens);
97        let mut inner = self.inner.lock().unwrap();
98        let eu = inner
99            .per_endpoint
100            .entry(endpoint_name.to_owned())
101            .or_default();
102        eu.calls += 1;
103        eu.input_tokens += input_tokens;
104        eu.output_tokens += output_tokens;
105        eu.cost += cost;
106    }
107
108    /// Records a flat-cost call (MCP tool, agent task).
109    ///
110    /// # Panics
111    ///
112    /// Panics if the mutex is poisoned.
113    #[allow(clippy::significant_drop_tightening)]
114    pub fn record_flat_call(&self, endpoint_name: &str, endpoint: &ResolvedEndpoint) {
115        let mut inner = self.inner.lock().unwrap();
116        let eu = inner
117            .per_endpoint
118            .entry(endpoint_name.to_owned())
119            .or_default();
120        eu.calls += 1;
121        eu.cost += endpoint.per_call_price;
122    }
123
124    /// Reads current usage without locking for the full snapshot.
125    ///
126    /// # Panics
127    ///
128    /// Panics if the mutex is poisoned.
129    #[must_use]
130    pub fn current_test_snapshot(&self) -> UsageSnapshot {
131        let inner = self.inner.lock().unwrap();
132        UsageSnapshot::from_endpoints(&inner.per_endpoint)
133    }
134
135    /// Reads the global aggregated snapshot.
136    ///
137    /// # Panics
138    ///
139    /// Panics if the mutex is poisoned.
140    #[must_use]
141    pub fn global_snapshot(&self) -> UsageSnapshot {
142        let inner = self.inner.lock().unwrap();
143        inner.global.clone()
144    }
145
146    /// Reads per-test snapshots.
147    ///
148    /// # Panics
149    ///
150    /// Panics if the mutex is poisoned.
151    #[must_use]
152    pub fn per_test_snapshots(&self) -> Vec<(String, UsageSnapshot)> {
153        let inner = self.inner.lock().unwrap();
154        inner.per_test.clone()
155    }
156
157    /// Resets the per-test accumulator. Call at the start of each test.
158    ///
159    /// # Panics
160    ///
161    /// Panics if the mutex is poisoned.
162    pub fn reset_per_test(&self) {
163        let mut inner = self.inner.lock().unwrap();
164        inner.per_endpoint.clear();
165    }
166
167    /// Commits the current test's usage to the global accumulator and stores
168    /// it as a per-test snapshot.
169    ///
170    /// # Panics
171    ///
172    /// Panics if the mutex is poisoned.
173    pub fn commit_test(&self, test_name: &str) {
174        let mut inner = self.inner.lock().unwrap();
175        let snapshot = UsageSnapshot::from_endpoints(&inner.per_endpoint);
176        // Merge into global
177        let ep_snapshot = inner.per_endpoint.clone();
178        for (ep_name, ep_usage) in &ep_snapshot {
179            let ge = inner.global.endpoints.entry(ep_name.clone()).or_default();
180            ge.calls += ep_usage.calls;
181            ge.input_tokens += ep_usage.input_tokens;
182            ge.output_tokens += ep_usage.output_tokens;
183            ge.cost += ep_usage.cost;
184        }
185        inner.global.total_cost += snapshot.total_cost;
186        inner.global.total_calls += snapshot.total_calls;
187        inner.global.total_tokens += snapshot.total_tokens;
188        inner.per_test.push((test_name.to_owned(), snapshot));
189    }
190}
191
192impl Default for UsageTracker {
193    fn default() -> Self {
194        Self::new()
195    }
196}
197
198/// Calculates the cost of an LLM call based on token pricing.
199#[allow(clippy::cast_precision_loss, clippy::suboptimal_flops)]
200#[must_use]
201pub fn calculate_llm_cost(
202    endpoint: &ResolvedEndpoint,
203    input_tokens: u64,
204    output_tokens: u64,
205) -> f64 {
206    let input_cost = (input_tokens as f64 / 1_000_000.0) * endpoint.input_price_per_1m;
207    let output_cost = (output_tokens as f64 / 1_000_000.0) * endpoint.output_price_per_1m;
208    input_cost + output_cost
209}
210
211/// Usage info extracted from an LLM API response.
212#[derive(Debug, Default, Clone, Copy)]
213pub struct LlmUsage {
214    /// Number of prompt / input tokens.
215    pub prompt_tokens: u64,
216    /// Number of completion / output tokens.
217    pub completion_tokens: u64,
218    /// Total tokens used.
219    pub total_tokens: u64,
220}
221
222/// Result of an LLM chat call including usage data.
223#[derive(Debug, Clone)]
224pub struct LlmResponse {
225    /// The message content from the LLM.
226    pub content: String,
227    /// Token usage from the API response.
228    pub usage: LlmUsage,
229}
230
231/// Extracts token usage from an OpenAI-compatible API response JSON.
232#[must_use]
233pub fn extract_usage(value: &serde_json::Value) -> LlmUsage {
234    let usage = &value["usage"];
235    LlmUsage {
236        prompt_tokens: usage["prompt_tokens"].as_u64().unwrap_or(0),
237        completion_tokens: usage["completion_tokens"].as_u64().unwrap_or(0),
238        total_tokens: usage["total_tokens"].as_u64().unwrap_or(0),
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use crate::costs::{calculate_llm_cost, UsageTracker};
245    use crate::endpoints::ResolvedEndpoint;
246    use crate::scenario::EndpointType;
247
248    fn make_endpoint(
249        name: &str,
250        input_price: f64,
251        output_price: f64,
252        per_call: f64,
253    ) -> ResolvedEndpoint {
254        ResolvedEndpoint {
255            name: name.to_owned(),
256            endpoint_type: EndpointType::Llm,
257            url: String::new(),
258            model: None,
259            api_key: None,
260            headers: std::collections::HashMap::new(),
261            command: None,
262            args: vec![],
263            input_price_per_1m: input_price,
264            output_price_per_1m: output_price,
265            per_call_price: per_call,
266        }
267    }
268
269    #[test]
270    fn test_calculate_llm_cost() {
271        let ep = make_endpoint("test", 0.15, 0.60, 0.0);
272        // 1M input tokens = $0.15, 500K output = $0.30
273        let cost = calculate_llm_cost(&ep, 1_000_000, 500_000);
274        assert!((cost - 0.45).abs() < 0.001);
275    }
276
277    #[test]
278    fn test_calculate_zero_cost() {
279        let ep = make_endpoint("free", 0.0, 0.0, 0.0);
280        let cost = calculate_llm_cost(&ep, 1_000_000, 1_000_000);
281        assert!((cost - 0.0).abs() < f64::EPSILON);
282    }
283
284    #[test]
285    fn test_usage_tracker_record_llm() {
286        let tracker = UsageTracker::new();
287        let ep = make_endpoint("gpt4", 2.50, 10.0, 0.0);
288        tracker.record_llm_call("gpt4", &ep, 1000, 500);
289
290        let snap = tracker.current_test_snapshot();
291        assert_eq!(snap.total_calls, 1);
292        assert_eq!(snap.total_tokens, 1500);
293        assert!(
294            snap.total_cost > 0.0,
295            "expected cost > 0, got {}",
296            snap.total_cost
297        );
298
299        let ep_usage = snap.endpoints.get("gpt4").unwrap();
300        assert_eq!(ep_usage.calls, 1);
301        assert_eq!(ep_usage.input_tokens, 1000);
302        assert_eq!(ep_usage.output_tokens, 500);
303    }
304
305    #[test]
306    fn test_usage_tracker_record_flat() {
307        let tracker = UsageTracker::new();
308        let ep = make_endpoint("agent", 0.0, 0.0, 0.01);
309        tracker.record_flat_call("agent", &ep);
310        tracker.record_flat_call("agent", &ep);
311
312        let snap = tracker.current_test_snapshot();
313        assert_eq!(snap.total_calls, 2);
314        assert!((snap.total_cost - 0.02).abs() < f64::EPSILON);
315    }
316
317    #[test]
318    fn test_usage_tracker_multiple_endpoints() {
319        let tracker = UsageTracker::new();
320        let fast = make_endpoint("fast", 0.15, 0.60, 0.0);
321        let slow = make_endpoint("slow", 2.50, 10.0, 0.0);
322
323        tracker.record_llm_call("fast", &fast, 100, 50);
324        tracker.record_llm_call("slow", &slow, 200, 100);
325
326        let snap = tracker.current_test_snapshot();
327        assert_eq!(snap.total_calls, 2);
328        assert_eq!(snap.endpoints.len(), 2);
329    }
330
331    #[test]
332    fn test_usage_tracker_reset_and_commit() {
333        let tracker = UsageTracker::new();
334        let ep = make_endpoint("test", 0.15, 0.60, 0.0);
335
336        tracker.record_llm_call("test", &ep, 100, 50);
337        tracker.commit_test("test1");
338        tracker.reset_per_test();
339
340        tracker.record_llm_call("test", &ep, 200, 100);
341        tracker.commit_test("test2");
342
343        let global = tracker.global_snapshot();
344        assert_eq!(global.total_calls, 2);
345        assert_eq!(global.total_tokens, 450);
346
347        let per_test = tracker.per_test_snapshots();
348        assert_eq!(per_test.len(), 2);
349        assert_eq!(per_test[0].0, "test1");
350        assert_eq!(per_test[1].0, "test2");
351    }
352}
llm_browser_testkit/costs.rs

llm_browser_testkit/
costs.rs
