lean_ctx/proxy/
usage.rs

1//! Real provider-reported token usage extraction.
2//!
3//! The proxy already rewrites requests; this module reads the *response* so the
4//! dashboard/terminal can show **measured** cost (the user's real provider bill)
5//! instead of an estimate. All three providers report the exact model and the
6//! billed token breakdown — including prompt-cache reads/writes — in the final
7//! event of a stream (or the body of a non-streaming response):
8//!
9//! - Anthropic: `message_start` carries model + input/cache tokens, the final
10//!   `message_delta` carries `output_tokens`. Non-streaming: one `usage` object.
11//! - OpenAI Responses: the `response.completed` event nests `response.usage`.
12//! - OpenAI Chat Completions: the final chunk carries `usage` (needs
13//!   `stream_options.include_usage`, which the proxy injects).
14//! - Gemini: every chunk carries `usageMetadata`; the last one has the totals.
15//!
16//! [`RealUsage`] normalizes every provider onto the four billable buckets that
17//! [`crate::core::gain::model_pricing::ModelCost::estimate_usd`] prices:
18//! uncached input, output (incl. reasoning/thoughts), cache-read, cache-write.
19
20use futures::{Stream, StreamExt};
21use serde_json::Value;
22
23/// LLM provider whose response shape a [`Scanner`] understands.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum Provider {
26    Anthropic,
27    /// Covers both Chat Completions and the Responses API (same `usage` dialects
28    /// are detected by field name).
29    OpenAi,
30    Gemini,
31}
32
33impl Provider {
34    /// Maps the proxy's `provider_label` (`"Anthropic"`, `"OpenAI"`/`"ChatGPT"`, else Gemini).
35    pub fn from_label(label: &str) -> Self {
36        match label {
37            "Anthropic" => Self::Anthropic,
38            "OpenAI" | "ChatGPT" => Self::OpenAi,
39            _ => Self::Gemini,
40        }
41    }
42}
43
44/// One LLM turn's real, provider-reported usage, normalized to billable buckets.
45///
46/// `output_tokens` already includes reasoning/thinking tokens (they are billed at
47/// the output rate); `reasoning_tokens` is retained only for display.
48#[derive(Debug, Clone, Default, PartialEq)]
49pub struct RealUsage {
50    pub model: String,
51    /// Input tokens billed at the input rate (cache reads/writes excluded).
52    pub input_tokens: u64,
53    /// Output tokens billed at the output rate (includes reasoning/thoughts).
54    pub output_tokens: u64,
55    /// Input tokens served from the prompt cache (billed at the cache-read rate).
56    pub cache_read_tokens: u64,
57    /// Input tokens written to the prompt cache (Anthropic only; 0 elsewhere).
58    pub cache_write_tokens: u64,
59    /// Reasoning/thinking subset of `output_tokens` (display only).
60    pub reasoning_tokens: u64,
61    /// Output-savings experiment arm for this turn (#895 Track B), or `None` when
62    /// no holdout is active. Stamped from the request, not parsed from the
63    /// response — it identifies whether this turn was output-shaped.
64    pub cohort: Option<super::holdout::Arm>,
65}
66
67impl RealUsage {
68    /// True once any model or token field has been observed — the gate for
69    /// recording. Avoids emitting empty rows for streams that never reported usage.
70    fn is_meaningful(&self) -> bool {
71        !self.model.is_empty()
72            || self.input_tokens > 0
73            || self.output_tokens > 0
74            || self.cache_read_tokens > 0
75            || self.cache_write_tokens > 0
76    }
77}
78
79/// Upper bound on a single buffered line before we give up on it. Usage events
80/// are tiny; this only guards against a pathological newline-free stream.
81const MAX_LINE_BYTES: usize = 1 << 20; // 1 MiB
82
83/// Incrementally extracts [`RealUsage`] from a response stream (or a full body).
84///
85/// `feed` is called with raw response chunks and keeps only the trailing partial
86/// line buffered (O(1) memory beyond one line); `finalize` returns the merged
87/// usage once the stream ends.
88pub struct Scanner {
89    provider: Provider,
90    /// Model parsed from the request URL (Gemini puts it there, not in the body).
91    url_model: Option<String>,
92    /// Output-savings arm (#895), stamped onto the usage at finalize.
93    cohort: Option<super::holdout::Arm>,
94    buf: Vec<u8>,
95    usage: RealUsage,
96}
97
98impl Scanner {
99    pub fn new(provider: Provider, url_model: Option<String>) -> Self {
100        Self {
101            provider,
102            url_model,
103            cohort: None,
104            buf: Vec::new(),
105            usage: RealUsage::default(),
106        }
107    }
108
109    /// Tags the usage this scanner produces with an output-savings arm (#895).
110    #[must_use]
111    pub fn with_cohort(mut self, cohort: Option<super::holdout::Arm>) -> Self {
112        self.cohort = cohort;
113        self
114    }
115
116    /// Feeds a raw streaming chunk, scanning every complete line it completes.
117    pub fn feed(&mut self, chunk: &[u8]) {
118        self.buf.extend_from_slice(chunk);
119        while let Some(nl) = self.buf.iter().position(|&b| b == b'\n') {
120            let mut line: Vec<u8> = self.buf.drain(..=nl).collect();
121            line.pop(); // drop '\n'
122            if line.last() == Some(&b'\r') {
123                line.pop();
124            }
125            self.scan_line(&line);
126        }
127        if self.buf.len() > MAX_LINE_BYTES {
128            self.buf.clear();
129        }
130    }
131
132    /// Feeds a complete non-streaming JSON response body.
133    pub fn feed_body(&mut self, body: &[u8]) {
134        if let Ok(v) = serde_json::from_slice::<Value>(body) {
135            self.absorb(&v);
136        }
137    }
138
139    /// Consumes the scanner, flushing any trailing partial line (a final event
140    /// may arrive without a newline) and returning the merged usage if any.
141    pub fn finalize(mut self) -> Option<RealUsage> {
142        if !self.buf.is_empty() {
143            let line = std::mem::take(&mut self.buf);
144            self.scan_line(&line);
145        }
146        if self.usage.is_meaningful() {
147            self.usage.cohort = self.cohort;
148            Some(self.usage)
149        } else {
150            None
151        }
152    }
153
154    fn scan_line(&mut self, line: &[u8]) {
155        let Ok(text) = std::str::from_utf8(line) else {
156            return;
157        };
158        let trimmed = text.trim();
159        if trimmed.is_empty() {
160            return;
161        }
162        // Cheap pre-filter: skip the bulk of the stream (content deltas) and only
163        // JSON-parse lines that can carry usage or the model name.
164        if !self.line_might_be_relevant(trimmed) {
165            return;
166        }
167        let json_str = if let Some(rest) = trimmed.strip_prefix("data:") {
168            // SSE (Anthropic, OpenAI, Gemini with alt=sse).
169            let r = rest.trim();
170            if r.is_empty() || r == "[DONE]" {
171                return;
172            }
173            r
174        } else if trimmed.starts_with('{') {
175            // NDJSON / array-element line (Gemini x-ndjson). Tolerate the array
176            // punctuation a streamed JSON array puts around an element.
177            trimmed
178                .trim_start_matches([',', '['])
179                .trim_end_matches([',', ']'])
180                .trim()
181        } else {
182            return;
183        };
184        if let Ok(v) = serde_json::from_str::<Value>(json_str) {
185            self.absorb(&v);
186        }
187    }
188
189    fn line_might_be_relevant(&self, s: &str) -> bool {
190        match self.provider {
191            // Anthropic `message_start`/`message_delta` and OpenAI `usage`/
192            // `response.*` events all contain the substring "usage".
193            Provider::Anthropic | Provider::OpenAi => s.contains("usage"),
194            Provider::Gemini => s.contains("usageMetadata"),
195        }
196    }
197
198    fn absorb(&mut self, v: &Value) {
199        match self.provider {
200            Provider::Anthropic => absorb_anthropic(&mut self.usage, v),
201            Provider::OpenAi => absorb_openai(&mut self.usage, v),
202            Provider::Gemini => absorb_gemini(&mut self.usage, v, self.url_model.as_deref()),
203        }
204    }
205}
206
207/// Anthropic: model + input/cache live on `message` (streaming `message_start`
208/// or a non-streaming body); `output_tokens` arrives later on the event-level
209/// `usage` of `message_delta`. Latest non-zero wins, so the cumulative final
210/// delta is authoritative.
211fn absorb_anthropic(u: &mut RealUsage, v: &Value) {
212    let msg = v.get("message").unwrap_or(v);
213    if let Some(model) = msg.get("model").and_then(Value::as_str)
214        && !model.is_empty()
215    {
216        u.model = model.to_string();
217    }
218    let Some(usage) = msg.get("usage").or_else(|| v.get("usage")) else {
219        return;
220    };
221    if let Some(n) = usage.get("input_tokens").and_then(Value::as_u64) {
222        u.input_tokens = n;
223    }
224    if let Some(n) = usage.get("cache_read_input_tokens").and_then(Value::as_u64) {
225        u.cache_read_tokens = n;
226    }
227    if let Some(n) = usage
228        .get("cache_creation_input_tokens")
229        .and_then(Value::as_u64)
230    {
231        u.cache_write_tokens = n;
232    }
233    if let Some(n) = usage.get("output_tokens").and_then(Value::as_u64)
234        && n > 0
235    {
236        u.output_tokens = n;
237    }
238}
239
240/// OpenAI Chat Completions + Responses. `response.completed` nests the payload
241/// under `response`; chat chunks and non-streaming bodies are top-level. Both
242/// `usage` dialects are accepted (Responses: `input_tokens`/`output_tokens`;
243/// Chat: `prompt_tokens`/`completion_tokens`). `cached_tokens` is the cache-read
244/// portion of the reported input; OpenAI bills no separate cache write.
245fn absorb_openai(u: &mut RealUsage, v: &Value) {
246    let root = v.get("response").unwrap_or(v);
247    if let Some(model) = root.get("model").and_then(Value::as_str)
248        && !model.is_empty()
249    {
250        u.model = model.to_string();
251    }
252    let Some(usage) = root.get("usage") else {
253        return;
254    };
255    if usage.is_null() {
256        // `response.created` / `response.in_progress` carry `usage: null`.
257        return;
258    }
259
260    let total_input = usage
261        .get("input_tokens")
262        .or_else(|| usage.get("prompt_tokens"))
263        .and_then(Value::as_u64)
264        .unwrap_or(0);
265    let total_output = usage
266        .get("output_tokens")
267        .or_else(|| usage.get("completion_tokens"))
268        .and_then(Value::as_u64)
269        .unwrap_or(0);
270    let cached = usage
271        .get("input_tokens_details")
272        .or_else(|| usage.get("prompt_tokens_details"))
273        .and_then(|d| d.get("cached_tokens"))
274        .and_then(Value::as_u64)
275        .unwrap_or(0);
276    let reasoning = usage
277        .get("output_tokens_details")
278        .or_else(|| usage.get("completion_tokens_details"))
279        .and_then(|d| d.get("reasoning_tokens"))
280        .and_then(Value::as_u64)
281        .unwrap_or(0);
282
283    if total_input == 0 && total_output == 0 {
284        return;
285    }
286    u.input_tokens = total_input.saturating_sub(cached);
287    u.cache_read_tokens = cached;
288    u.cache_write_tokens = 0;
289    u.output_tokens = total_output;
290    u.reasoning_tokens = reasoning;
291}
292
293/// Gemini: `usageMetadata` carries the counts; `modelVersion` (or the request
294/// URL) the model. `thoughtsTokenCount` is billed at the output rate and is not
295/// part of `candidatesTokenCount`, so it is added into `output_tokens`.
296fn absorb_gemini(u: &mut RealUsage, v: &Value, url_model: Option<&str>) {
297    if let Some(mv) = v.get("modelVersion").and_then(Value::as_str)
298        && !mv.is_empty()
299    {
300        u.model = mv.to_string();
301    } else if u.model.is_empty()
302        && let Some(m) = url_model
303        && !m.is_empty()
304    {
305        u.model = m.to_string();
306    }
307    let Some(um) = v.get("usageMetadata") else {
308        return;
309    };
310    let prompt = um
311        .get("promptTokenCount")
312        .and_then(Value::as_u64)
313        .unwrap_or(0);
314    let candidates = um
315        .get("candidatesTokenCount")
316        .and_then(Value::as_u64)
317        .unwrap_or(0);
318    let cached = um
319        .get("cachedContentTokenCount")
320        .and_then(Value::as_u64)
321        .unwrap_or(0);
322    let thoughts = um
323        .get("thoughtsTokenCount")
324        .and_then(Value::as_u64)
325        .unwrap_or(0);
326    if prompt == 0 && candidates == 0 && thoughts == 0 {
327        return;
328    }
329    u.input_tokens = prompt.saturating_sub(cached);
330    u.cache_read_tokens = cached;
331    u.cache_write_tokens = 0;
332    u.output_tokens = candidates + thoughts;
333    u.reasoning_tokens = thoughts;
334}
335
336/// Extracts the model from a Gemini request path
337/// (`/v1beta/models/{model}:generateContent`). Returns `None` for other paths.
338pub fn gemini_model_from_path(path: &str) -> Option<String> {
339    let after = path.rsplit_once("/models/").map(|(_, m)| m)?;
340    let model = after.split(':').next().unwrap_or(after).trim();
341    if model.is_empty() {
342        None
343    } else {
344        Some(model.to_string())
345    }
346}
347
348/// Wraps a response byte stream so every chunk is forwarded **byte-for-byte**
349/// while a [`Scanner`] observes it; on stream end the merged usage is recorded.
350/// Memory overhead is one buffered SSE line, never the whole response.
351pub fn tee_stream<S, B, E>(
352    inner: S,
353    scanner: Scanner,
354) -> impl Stream<Item = Result<B, E>> + Send + 'static
355where
356    S: Stream<Item = Result<B, E>> + Send + Unpin + 'static,
357    B: AsRef<[u8]> + Send + 'static,
358    E: Send + 'static,
359{
360    futures::stream::unfold(
361        (inner, Some(scanner)),
362        |(mut inner, mut scanner)| async move {
363            match inner.next().await {
364                Some(Ok(chunk)) => {
365                    if let Some(s) = scanner.as_mut() {
366                        s.feed(chunk.as_ref());
367                    }
368                    Some((Ok(chunk), (inner, scanner)))
369                }
370                Some(err) => Some((err, (inner, scanner))),
371                None => {
372                    if let Some(s) = scanner.take()
373                        && let Some(usage) = s.finalize()
374                    {
375                        super::usage_meter::record(&usage);
376                    }
377                    None
378                }
379            }
380        },
381    )
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387
388    fn feed_lines(
389        provider: Provider,
390        url_model: Option<&str>,
391        lines: &[&str],
392    ) -> Option<RealUsage> {
393        let mut s = Scanner::new(provider, url_model.map(str::to_string));
394        for line in lines {
395            s.feed(line.as_bytes());
396            s.feed(b"\n");
397        }
398        s.finalize()
399    }
400
401    #[test]
402    fn anthropic_merges_message_start_and_delta() {
403        let u = feed_lines(
404            Provider::Anthropic,
405            None,
406            &[
407                r#"data: {"type":"message_start","message":{"model":"claude-opus-4-5-20251101","usage":{"input_tokens":100,"cache_read_input_tokens":2000,"cache_creation_input_tokens":50,"output_tokens":1}}}"#,
408                r#"data: {"type":"content_block_delta","index":0,"delta":{"text":"hello"}}"#,
409                r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":73}}"#,
410                "data: {\"type\":\"message_stop\"}",
411            ],
412        )
413        .expect("usage");
414        assert_eq!(u.model, "claude-opus-4-5-20251101");
415        assert_eq!(u.input_tokens, 100);
416        assert_eq!(u.cache_read_tokens, 2000);
417        assert_eq!(u.cache_write_tokens, 50);
418        assert_eq!(u.output_tokens, 73);
419    }
420
421    #[test]
422    fn anthropic_non_streaming_body() {
423        let mut s = Scanner::new(Provider::Anthropic, None);
424        s.feed_body(
425            br#"{"model":"claude-sonnet-4-5","usage":{"input_tokens":24,"output_tokens":18,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}"#,
426        );
427        let u = s.finalize().expect("usage");
428        assert_eq!(u.model, "claude-sonnet-4-5");
429        assert_eq!(u.input_tokens, 24);
430        assert_eq!(u.output_tokens, 18);
431    }
432
433    #[test]
434    fn openai_responses_completed_event() {
435        let u = feed_lines(
436            Provider::OpenAi,
437            None,
438            &[
439                r#"data: {"type":"response.created","response":{"model":"gpt-5.4","usage":null}}"#,
440                r#"data: {"type":"response.completed","response":{"model":"gpt-5.4","usage":{"input_tokens":1289,"input_tokens_details":{"cached_tokens":289},"output_tokens":685,"output_tokens_details":{"reasoning_tokens":640},"total_tokens":1974}}}"#,
441            ],
442        )
443        .expect("usage");
444        assert_eq!(u.model, "gpt-5.4");
445        assert_eq!(u.input_tokens, 1000); // 1289 - 289 cached
446        assert_eq!(u.cache_read_tokens, 289);
447        assert_eq!(u.output_tokens, 685);
448        assert_eq!(u.reasoning_tokens, 640);
449    }
450
451    #[test]
452    fn openai_chat_final_usage_chunk() {
453        let u = feed_lines(
454            Provider::OpenAi,
455            None,
456            &[
457                r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"gpt-5.4-mini"}"#,
458                r#"data: {"choices":[],"model":"gpt-5.4-mini","usage":{"prompt_tokens":500,"prompt_tokens_details":{"cached_tokens":100},"completion_tokens":40,"total_tokens":540}}"#,
459                "data: [DONE]",
460            ],
461        )
462        .expect("usage");
463        assert_eq!(u.model, "gpt-5.4-mini");
464        assert_eq!(u.input_tokens, 400);
465        assert_eq!(u.cache_read_tokens, 100);
466        assert_eq!(u.output_tokens, 40);
467    }
468
469    #[test]
470    fn gemini_usage_metadata_with_url_model() {
471        let u = feed_lines(
472            Provider::Gemini,
473            Some("gemini-2.5-pro"),
474            &[
475                r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":25,"candidatesTokenCount":7,"thoughtsTokenCount":39,"totalTokenCount":71}}"#,
476            ],
477        )
478        .expect("usage");
479        assert_eq!(u.model, "gemini-2.5-pro");
480        assert_eq!(u.input_tokens, 25);
481        assert_eq!(u.output_tokens, 46); // candidates 7 + thoughts 39
482        assert_eq!(u.reasoning_tokens, 39);
483    }
484
485    #[test]
486    fn gemini_prefers_model_version_over_url() {
487        let u = feed_lines(
488            Provider::Gemini,
489            Some("gemini-2.5-pro"),
490            &[
491                r#"data: {"modelVersion":"gemini-2.5-pro-002","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"cachedContentTokenCount":4,"totalTokenCount":15}}"#,
492            ],
493        )
494        .expect("usage");
495        assert_eq!(u.model, "gemini-2.5-pro-002");
496        assert_eq!(u.input_tokens, 6); // 10 - 4 cached
497        assert_eq!(u.cache_read_tokens, 4);
498        assert_eq!(u.output_tokens, 5);
499    }
500
501    #[test]
502    fn split_chunks_reassemble_across_feed_calls() {
503        let mut s = Scanner::new(Provider::Anthropic, None);
504        let line = r#"data: {"type":"message_delta","delta":{},"usage":{"output_tokens":42}}"#;
505        let bytes = format!("{line}\n");
506        let (a, b) = bytes.as_bytes().split_at(20);
507        s.feed(a);
508        s.feed(b);
509        let u = s.finalize().expect("usage");
510        assert_eq!(u.output_tokens, 42);
511    }
512
513    #[test]
514    fn no_usage_yields_none() {
515        let out = feed_lines(
516            Provider::OpenAi,
517            None,
518            &[r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"gpt-5.4"}"#],
519        );
520        assert!(out.is_none(), "content-only stream reports no usage");
521    }
522
523    #[test]
524    fn final_event_without_trailing_newline() {
525        let mut s = Scanner::new(Provider::Anthropic, None);
526        s.feed(br#"data: {"type":"message_delta","delta":{},"usage":{"output_tokens":7}}"#);
527        let u = s.finalize().expect("flushes trailing partial line");
528        assert_eq!(u.output_tokens, 7);
529    }
530
531    #[test]
532    fn gemini_model_from_path_extracts() {
533        assert_eq!(
534            gemini_model_from_path("/v1beta/models/gemini-2.5-pro:streamGenerateContent")
535                .as_deref(),
536            Some("gemini-2.5-pro")
537        );
538        assert_eq!(
539            gemini_model_from_path("/v1beta/models/gemini-2.5-flash:generateContent").as_deref(),
540            Some("gemini-2.5-flash")
541        );
542        assert_eq!(gemini_model_from_path("/v1/chat/completions"), None);
543    }
544
545    #[tokio::test]
546    async fn tee_stream_passes_bytes_through_and_records() {
547        let chunks: Vec<Result<Vec<u8>, std::convert::Infallible>> = vec![
548            Ok(b"data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-haiku-4.5\",\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}}\n".to_vec()),
549            Ok(b"data: {\"type\":\"message_delta\",\"delta\":{},\"usage\":{\"output_tokens\":9}}\n".to_vec()),
550        ];
551        let inner = futures::stream::iter(chunks);
552        let scanner = Scanner::new(Provider::Anthropic, None);
553        let teed = tee_stream(inner, scanner);
554        let collected: Vec<_> = teed.collect().await;
555        // Byte-for-byte passthrough preserved.
556        assert_eq!(collected.len(), 2);
557        assert!(collected.iter().all(std::result::Result::is_ok));
558    }
559}
lean_ctx/proxy/usage.rs

lean_ctx/proxy/
usage.rs
