Skip to main content

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 + OpenRouter models
58    /// with explicit cache-write pricing; 0 elsewhere).
59    pub cache_write_tokens: u64,
60    /// Reasoning/thinking subset of `output_tokens` (display only).
61    pub reasoning_tokens: u64,
62    /// USD the provider *actually charged* for this turn, when the response
63    /// reports it. OpenRouter: `usage.cost` in credits (≡ USD); for BYOK
64    /// requests the separate `cost_details.upstream_inference_cost` is added
65    /// when it differs from `cost` (#746 — non-BYOK mirrors cost there).
66    /// The measured figure beats any price-table estimate wherever both exist.
67    /// `None` for providers that report tokens only (Anthropic/OpenAI/Gemini).
68    pub provider_cost_usd: Option<f64>,
69    /// Output-savings experiment arm for this turn (#895 Track B), or `None` when
70    /// no holdout is active. Stamped from the request, not parsed from the
71    /// response — it identifies whether this turn was output-shaped.
72    pub cohort: Option<super::holdout::Arm>,
73    /// Request-side gateway context (enterprise#11/#17/#18): identity tags,
74    /// compression savings and baseline inputs, stamped from the request before
75    /// it left for the upstream. `None` outside the forward path (e.g. tests
76    /// that only parse response bodies).
77    pub wire: Option<Box<WireContext>>,
78}
79
80/// Request-side context the forward path knows and the response scanner does
81/// not: who sent the request (gateway identity, enterprise#11), what the proxy
82/// saved on the wire, and the counterfactual-baseline inputs (enterprise#18).
83/// Travels inside [`RealUsage`] so `usage_meter::record` stays the single
84/// choke-point through which every measured turn flows.
85#[derive(Debug, Clone, Default, PartialEq)]
86pub struct WireContext {
87    /// Provider label (`Anthropic|OpenAI|ChatGPT|Gemini`) of the serving route.
88    pub provider: String,
89    /// Person tag from the gateway key (enterprise#11).
90    pub person: Option<String>,
91    /// Team tag from the gateway key.
92    pub team: Option<String>,
93    /// Project: `x-leanctx-project` header wins over the key's default project.
94    pub project: Option<String>,
95    /// Estimated tokens this request saved through wire compression
96    /// (bytes/4 heuristic, same basis as the proxy stats).
97    pub saved_tokens: u64,
98    /// Estimated request tokens BEFORE lean-ctx compression (bytes/4) — the
99    /// SEE-attribution input of the avoided-cost baseline (enterprise#18).
100    pub uncompressed_input_tokens: u64,
101    /// True when the serving upstream is a local/loopback endpoint — billed via
102    /// the transparent `local_shadow_rate`, never 0 (enterprise#15/#18).
103    pub is_local: bool,
104    /// Originally requested model when the router downgraded/aliased it
105    /// (enterprise#13); `None` for passthrough.
106    pub routed_from: Option<String>,
107    /// Provider-counted input tokens of the ORIGINAL (uncompressed) request,
108    /// from the free `count_tokens` probe (#701). `None` when metering is off
109    /// or no probe was spawned; an empty slot at read time (probe failed or
110    /// still in flight) degrades the row to the local estimate.
111    pub counterfactual: Option<super::counterfactual::CounterfactualSlot>,
112}
113
114impl RealUsage {
115    /// True once any model, token or measured-cost field has been observed —
116    /// the gate for recording. Avoids emitting empty rows for streams that
117    /// never reported usage.
118    fn is_meaningful(&self) -> bool {
119        !self.model.is_empty()
120            || self.input_tokens > 0
121            || self.output_tokens > 0
122            || self.cache_read_tokens > 0
123            || self.cache_write_tokens > 0
124            || self.provider_cost_usd.is_some()
125    }
126}
127
128/// Upper bound on a single buffered line before we give up on it. Usage events
129/// are tiny; this only guards against a pathological newline-free stream.
130const MAX_LINE_BYTES: usize = 1 << 20; // 1 MiB
131
132/// Incrementally extracts [`RealUsage`] from a response stream (or a full body).
133///
134/// `feed` is called with raw response chunks and keeps only the trailing partial
135/// line buffered (O(1) memory beyond one line); `finalize` returns the merged
136/// usage once the stream ends.
137pub struct Scanner {
138    provider: Provider,
139    /// Model parsed from the request URL (Gemini puts it there, not in the body).
140    url_model: Option<String>,
141    /// Output-savings arm (#895), stamped onto the usage at finalize.
142    cohort: Option<super::holdout::Arm>,
143    /// Request-side gateway context (enterprise#11/#18), stamped at finalize.
144    wire: Option<Box<WireContext>>,
145    /// Billed USD from a gateway response header (#1189), stamped at finalize
146    /// unless the body already reported the charge.
147    header_cost: Option<f64>,
148    buf: Vec<u8>,
149    usage: RealUsage,
150}
151
152impl Scanner {
153    pub fn new(provider: Provider, url_model: Option<String>) -> Self {
154        Self {
155            provider,
156            url_model,
157            cohort: None,
158            wire: None,
159            header_cost: None,
160            buf: Vec::new(),
161            usage: RealUsage::default(),
162        }
163    }
164
165    /// Tags the usage this scanner produces with an output-savings arm (#895).
166    #[must_use]
167    pub fn with_cohort(mut self, cohort: Option<super::holdout::Arm>) -> Self {
168        self.cohort = cohort;
169        self
170    }
171
172    /// Attaches the request-side gateway context (identity tags, wire savings,
173    /// baseline inputs — enterprise#11/#17/#18) stamped onto the usage record.
174    #[must_use]
175    pub fn with_wire_context(mut self, wire: Option<Box<WireContext>>) -> Self {
176        self.wire = wire;
177        self
178    }
179
180    /// Attaches a billed USD figure reported by the upstream via a response
181    /// header (LiteLLM `x-litellm-response-cost`, or the operator-configured
182    /// `[proxy] cost_response_header`, #1189). Applied at finalize only when
183    /// the body did not already carry a measured cost — a body figure
184    /// (OpenRouter `usage.cost`) is the bill itself and always wins.
185    #[must_use]
186    pub fn with_header_cost(mut self, cost: Option<f64>) -> Self {
187        self.header_cost = cost.filter(|c| c.is_finite() && *c >= 0.0);
188        self
189    }
190
191    /// Feeds a raw streaming chunk, scanning every complete line it completes.
192    pub fn feed(&mut self, chunk: &[u8]) {
193        self.buf.extend_from_slice(chunk);
194        while let Some(nl) = self.buf.iter().position(|&b| b == b'\n') {
195            let mut line: Vec<u8> = self.buf.drain(..=nl).collect();
196            line.pop(); // drop '\n'
197            if line.last() == Some(&b'\r') {
198                line.pop();
199            }
200            self.scan_line(&line);
201        }
202        if self.buf.len() > MAX_LINE_BYTES {
203            self.buf.clear();
204        }
205    }
206
207    /// Feeds a complete non-streaming JSON response body.
208    pub fn feed_body(&mut self, body: &[u8]) {
209        if let Ok(v) = serde_json::from_slice::<Value>(body) {
210            self.absorb(&v);
211        }
212    }
213
214    /// Consumes the scanner, flushing any trailing partial line (a final event
215    /// may arrive without a newline) and returning the merged usage if any.
216    pub fn finalize(mut self) -> Option<RealUsage> {
217        if !self.buf.is_empty() {
218            let line = std::mem::take(&mut self.buf);
219            self.scan_line(&line);
220        }
221        // Gateway header cost (#1189): measured, but the body figure is the
222        // bill itself (OpenRouter usage.cost) and keeps priority when present.
223        if self.usage.provider_cost_usd.is_none() {
224            self.usage.provider_cost_usd = self.header_cost;
225        }
226        if self.usage.is_meaningful() {
227            self.usage.cohort = self.cohort;
228            self.usage.wire = self.wire;
229            Some(self.usage)
230        } else {
231            None
232        }
233    }
234
235    fn scan_line(&mut self, line: &[u8]) {
236        let Ok(text) = std::str::from_utf8(line) else {
237            return;
238        };
239        let trimmed = text.trim();
240        if trimmed.is_empty() {
241            return;
242        }
243        // Cheap pre-filter: skip the bulk of the stream (content deltas) and only
244        // JSON-parse lines that can carry usage or the model name.
245        if !self.line_might_be_relevant(trimmed) {
246            return;
247        }
248        let json_str = if let Some(rest) = trimmed.strip_prefix("data:") {
249            // SSE (Anthropic, OpenAI, Gemini with alt=sse).
250            let r = rest.trim();
251            if r.is_empty() || r == "[DONE]" {
252                return;
253            }
254            r
255        } else if trimmed.starts_with('{') {
256            // NDJSON / array-element line (Gemini x-ndjson). Tolerate the array
257            // punctuation a streamed JSON array puts around an element.
258            trimmed
259                .trim_start_matches([',', '['])
260                .trim_end_matches([',', ']'])
261                .trim()
262        } else {
263            return;
264        };
265        if let Ok(v) = serde_json::from_str::<Value>(json_str) {
266            self.absorb(&v);
267        }
268    }
269
270    fn line_might_be_relevant(&self, s: &str) -> bool {
271        match self.provider {
272            // Anthropic `message_start`/`message_delta` and OpenAI `usage`/
273            // `response.*` events all contain the substring "usage".
274            Provider::Anthropic | Provider::OpenAi => s.contains("usage"),
275            Provider::Gemini => s.contains("usageMetadata"),
276        }
277    }
278
279    fn absorb(&mut self, v: &Value) {
280        match self.provider {
281            Provider::Anthropic => absorb_anthropic(&mut self.usage, v),
282            Provider::OpenAi => absorb_openai(&mut self.usage, v),
283            Provider::Gemini => absorb_gemini(&mut self.usage, v, self.url_model.as_deref()),
284        }
285    }
286}
287
288/// Anthropic: model + input/cache live on `message` (streaming `message_start`
289/// or a non-streaming body); `output_tokens` arrives later on the event-level
290/// `usage` of `message_delta`. Latest non-zero wins, so the cumulative final
291/// delta is authoritative.
292fn absorb_anthropic(u: &mut RealUsage, v: &Value) {
293    let msg = v.get("message").unwrap_or(v);
294    if let Some(model) = msg.get("model").and_then(Value::as_str)
295        && !model.is_empty()
296    {
297        u.model = model.to_string();
298    }
299    let Some(usage) = msg.get("usage").or_else(|| v.get("usage")) else {
300        return;
301    };
302    if let Some(n) = usage.get("input_tokens").and_then(Value::as_u64) {
303        u.input_tokens = n;
304    }
305    if let Some(n) = usage.get("cache_read_input_tokens").and_then(Value::as_u64) {
306        u.cache_read_tokens = n;
307    }
308    if let Some(n) = usage
309        .get("cache_creation_input_tokens")
310        .and_then(Value::as_u64)
311    {
312        u.cache_write_tokens = n;
313    }
314    if let Some(n) = usage.get("output_tokens").and_then(Value::as_u64)
315        && n > 0
316    {
317        u.output_tokens = n;
318    }
319}
320
321/// OpenAI Chat Completions + Responses. `response.completed` nests the payload
322/// under `response`; chat chunks and non-streaming bodies are top-level. Both
323/// `usage` dialects are accepted (Responses: `input_tokens`/`output_tokens`;
324/// Chat: `prompt_tokens`/`completion_tokens`). `cached_tokens` is the cache-read
325/// portion of the reported input; OpenAI bills no separate cache write, but
326/// OpenRouter reports one (`prompt_tokens_details.cache_write_tokens`) for
327/// models with explicit cache-write pricing.
328///
329/// OpenRouter usage accounting additionally carries the money actually charged:
330/// `usage.cost` (credits ≡ USD) and, for BYOK requests, the upstream provider's
331/// own bill under `cost_details.upstream_inference_cost`. Their sum is this
332/// turn's real price — measured, not table-derived.
333fn absorb_openai(u: &mut RealUsage, v: &Value) {
334    let root = v.get("response").unwrap_or(v);
335    if let Some(model) = root.get("model").and_then(Value::as_str)
336        && !model.is_empty()
337    {
338        u.model = model.to_string();
339    }
340    let Some(usage) = root.get("usage") else {
341        return;
342    };
343    if usage.is_null() {
344        // `response.created` / `response.in_progress` carry `usage: null`.
345        return;
346    }
347
348    // Measured cost (OpenRouter dialect). Parsed before the token guard so a
349    // cost-bearing usage object is never lost, and `0` is preserved — a
350    // `:free` model's real price IS zero, not "unknown".
351    if let Some(cost) = usage.get("cost").and_then(Value::as_f64) {
352        let upstream = usage
353            .get("cost_details")
354            .and_then(|d| d.get("upstream_inference_cost"))
355            .and_then(Value::as_f64)
356            .unwrap_or(0.0);
357        // #746: non-BYOK responses may mirror `cost` in upstream_inference_cost;
358        // summing both would double-count. BYOK responses split the total:
359        // `cost` = OpenRouter fee (small), `upstream` = provider bill (large,
360        // always different from cost). Only add when genuinely distinct.
361        let byok_upstream = if upstream > 0.0 && upstream != cost {
362            upstream
363        } else {
364            0.0
365        };
366        u.provider_cost_usd = Some(cost + byok_upstream);
367    }
368
369    let total_input = usage
370        .get("input_tokens")
371        .or_else(|| usage.get("prompt_tokens"))
372        .and_then(Value::as_u64)
373        .unwrap_or(0);
374    let total_output = usage
375        .get("output_tokens")
376        .or_else(|| usage.get("completion_tokens"))
377        .and_then(Value::as_u64)
378        .unwrap_or(0);
379    let input_details = usage
380        .get("input_tokens_details")
381        .or_else(|| usage.get("prompt_tokens_details"));
382    let cached = input_details
383        .and_then(|d| d.get("cached_tokens"))
384        .and_then(Value::as_u64)
385        .unwrap_or(0);
386    let cache_write = input_details
387        .and_then(|d| d.get("cache_write_tokens"))
388        .and_then(Value::as_u64)
389        .unwrap_or(0);
390    let reasoning = usage
391        .get("output_tokens_details")
392        .or_else(|| usage.get("completion_tokens_details"))
393        .and_then(|d| d.get("reasoning_tokens"))
394        .and_then(Value::as_u64)
395        .unwrap_or(0);
396
397    if total_input == 0 && total_output == 0 {
398        return;
399    }
400    // OpenRouter counts cache writes inside prompt_tokens (unlike Anthropic's
401    // separate bucket) — subtract both cached reads and writes so the three
402    // buckets stay disjoint and are never double-priced.
403    u.input_tokens = total_input
404        .saturating_sub(cached)
405        .saturating_sub(cache_write);
406    u.cache_read_tokens = cached;
407    u.cache_write_tokens = cache_write;
408    u.output_tokens = total_output;
409    u.reasoning_tokens = reasoning;
410}
411
412/// Gemini: `usageMetadata` carries the counts; `modelVersion` (or the request
413/// URL) the model. `thoughtsTokenCount` is billed at the output rate and is not
414/// part of `candidatesTokenCount`, so it is added into `output_tokens`.
415fn absorb_gemini(u: &mut RealUsage, v: &Value, url_model: Option<&str>) {
416    if let Some(mv) = v.get("modelVersion").and_then(Value::as_str)
417        && !mv.is_empty()
418    {
419        u.model = mv.to_string();
420    } else if u.model.is_empty()
421        && let Some(m) = url_model
422        && !m.is_empty()
423    {
424        u.model = m.to_string();
425    }
426    let Some(um) = v.get("usageMetadata") else {
427        return;
428    };
429    let prompt = um
430        .get("promptTokenCount")
431        .and_then(Value::as_u64)
432        .unwrap_or(0);
433    let candidates = um
434        .get("candidatesTokenCount")
435        .and_then(Value::as_u64)
436        .unwrap_or(0);
437    let cached = um
438        .get("cachedContentTokenCount")
439        .and_then(Value::as_u64)
440        .unwrap_or(0);
441    let thoughts = um
442        .get("thoughtsTokenCount")
443        .and_then(Value::as_u64)
444        .unwrap_or(0);
445    if prompt == 0 && candidates == 0 && thoughts == 0 {
446        return;
447    }
448    u.input_tokens = prompt.saturating_sub(cached);
449    u.cache_read_tokens = cached;
450    u.cache_write_tokens = 0;
451    u.output_tokens = candidates + thoughts;
452    u.reasoning_tokens = thoughts;
453}
454
455/// Extracts the model from a Gemini request path
456/// (`/v1beta/models/{model}:generateContent`). Returns `None` for other paths.
457pub fn gemini_model_from_path(path: &str) -> Option<String> {
458    let after = path.rsplit_once("/models/").map(|(_, m)| m)?;
459    let model = after.split(':').next().unwrap_or(after).trim();
460    if model.is_empty() {
461        None
462    } else {
463        Some(model.to_string())
464    }
465}
466
467/// Wraps a response byte stream so every chunk is forwarded **byte-for-byte**
468/// while a [`Scanner`] observes it; on stream end the merged usage is recorded.
469/// Memory overhead is one buffered SSE line, never the whole response.
470pub fn tee_stream<S, B, E>(
471    inner: S,
472    scanner: Scanner,
473) -> impl Stream<Item = Result<B, E>> + Send + 'static
474where
475    S: Stream<Item = Result<B, E>> + Send + Unpin + 'static,
476    B: AsRef<[u8]> + Send + 'static,
477    E: Send + 'static,
478{
479    futures::stream::unfold(
480        (inner, Some(scanner)),
481        |(mut inner, mut scanner)| async move {
482            match inner.next().await {
483                Some(Ok(chunk)) => {
484                    if let Some(s) = scanner.as_mut() {
485                        s.feed(chunk.as_ref());
486                    }
487                    Some((Ok(chunk), (inner, scanner)))
488                }
489                Some(err) => Some((err, (inner, scanner))),
490                None => {
491                    if let Some(s) = scanner.take()
492                        && let Some(usage) = s.finalize()
493                    {
494                        super::usage_meter::record(&usage);
495                    }
496                    None
497                }
498            }
499        },
500    )
501}
502
503#[cfg(test)]
504mod tests {
505    use super::*;
506
507    fn feed_lines(
508        provider: Provider,
509        url_model: Option<&str>,
510        lines: &[&str],
511    ) -> Option<RealUsage> {
512        let mut s = Scanner::new(provider, url_model.map(str::to_string));
513        for line in lines {
514            s.feed(line.as_bytes());
515            s.feed(b"\n");
516        }
517        s.finalize()
518    }
519
520    #[test]
521    fn anthropic_merges_message_start_and_delta() {
522        let u = feed_lines(
523            Provider::Anthropic,
524            None,
525            &[
526                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}}}"#,
527                r#"data: {"type":"content_block_delta","index":0,"delta":{"text":"hello"}}"#,
528                r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":73}}"#,
529                "data: {\"type\":\"message_stop\"}",
530            ],
531        )
532        .expect("usage");
533        assert_eq!(u.model, "claude-opus-4-5-20251101");
534        assert_eq!(u.input_tokens, 100);
535        assert_eq!(u.cache_read_tokens, 2000);
536        assert_eq!(u.cache_write_tokens, 50);
537        assert_eq!(u.output_tokens, 73);
538    }
539
540    #[test]
541    fn anthropic_non_streaming_body() {
542        let mut s = Scanner::new(Provider::Anthropic, None);
543        s.feed_body(
544            br#"{"model":"claude-sonnet-4-5","usage":{"input_tokens":24,"output_tokens":18,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}"#,
545        );
546        let u = s.finalize().expect("usage");
547        assert_eq!(u.model, "claude-sonnet-4-5");
548        assert_eq!(u.input_tokens, 24);
549        assert_eq!(u.output_tokens, 18);
550    }
551
552    #[test]
553    fn scanner_stamps_wire_context_onto_usage() {
554        // enterprise#11/#17: the request-side context must survive scanning and
555        // arrive on the finalized record that usage_meter/store consume.
556        let wire = Box::new(WireContext {
557            provider: "Anthropic".into(),
558            person: Some("yves".into()),
559            team: None,
560            project: Some("billing".into()),
561            saved_tokens: 42,
562            uncompressed_input_tokens: 500,
563            is_local: false,
564            routed_from: None,
565            counterfactual: None,
566        });
567        let mut s = Scanner::new(Provider::Anthropic, None).with_wire_context(Some(wire.clone()));
568        s.feed_body(
569            br#"{"model":"claude-sonnet-4-5","usage":{"input_tokens":24,"output_tokens":18}}"#,
570        );
571        let u = s.finalize().expect("usage");
572        assert_eq!(u.wire, Some(wire));
573    }
574
575    #[test]
576    fn openai_responses_completed_event() {
577        let u = feed_lines(
578            Provider::OpenAi,
579            None,
580            &[
581                r#"data: {"type":"response.created","response":{"model":"gpt-5.4","usage":null}}"#,
582                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}}}"#,
583            ],
584        )
585        .expect("usage");
586        assert_eq!(u.model, "gpt-5.4");
587        assert_eq!(u.input_tokens, 1000); // 1289 - 289 cached
588        assert_eq!(u.cache_read_tokens, 289);
589        assert_eq!(u.output_tokens, 685);
590        assert_eq!(u.reasoning_tokens, 640);
591    }
592
593    #[test]
594    fn openai_chat_final_usage_chunk() {
595        let u = feed_lines(
596            Provider::OpenAi,
597            None,
598            &[
599                r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"gpt-5.4-mini"}"#,
600                r#"data: {"choices":[],"model":"gpt-5.4-mini","usage":{"prompt_tokens":500,"prompt_tokens_details":{"cached_tokens":100},"completion_tokens":40,"total_tokens":540}}"#,
601                "data: [DONE]",
602            ],
603        )
604        .expect("usage");
605        assert_eq!(u.model, "gpt-5.4-mini");
606        assert_eq!(u.input_tokens, 400);
607        assert_eq!(u.cache_read_tokens, 100);
608        assert_eq!(u.output_tokens, 40);
609        assert_eq!(u.provider_cost_usd, None, "OpenAI reports no usage.cost");
610    }
611
612    #[test]
613    fn openrouter_cost_and_cache_writes_are_measured() {
614        // #1179: OpenRouter usage accounting — the final streamed chunk carries
615        // the billed USD (`cost`) and the cache-write bucket inside
616        // prompt_tokens_details. All three token buckets must stay disjoint.
617        let u = feed_lines(
618            Provider::OpenAi,
619            None,
620            &[
621                r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"deepseek/deepseek-v4-flash-20260423"}"#,
622                r#"data: {"choices":[],"model":"deepseek/deepseek-v4-flash-20260423","usage":{"prompt_tokens":700,"prompt_tokens_details":{"cached_tokens":150,"cache_write_tokens":50},"completion_tokens":40,"cost":0.0123,"cost_details":{"upstream_inference_cost":null},"total_tokens":740}}"#,
623                "data: [DONE]",
624            ],
625        )
626        .expect("usage");
627        assert_eq!(u.input_tokens, 500, "700 - 150 cached - 50 cache-write");
628        assert_eq!(u.cache_read_tokens, 150);
629        assert_eq!(u.cache_write_tokens, 50);
630        assert_eq!(u.output_tokens, 40);
631        let cost = u.provider_cost_usd.expect("measured cost");
632        assert!((cost - 0.0123).abs() < 1e-12);
633    }
634
635    #[test]
636    fn openrouter_byok_adds_upstream_inference_cost() {
637        let mut s = Scanner::new(Provider::OpenAi, None);
638        s.feed_body(
639            br#"{"model":"anthropic/claude-sonnet-5","usage":{"prompt_tokens":100,"completion_tokens":10,"cost":0.05,"cost_details":{"upstream_inference_cost":0.95}}}"#,
640        );
641        let u = s.finalize().expect("usage");
642        let cost = u.provider_cost_usd.expect("measured cost");
643        assert!(
644            (cost - 1.0).abs() < 1e-12,
645            "OpenRouter fee + BYOK upstream bill"
646        );
647    }
648
649    /// #746: non-BYOK responses mirror `cost` in `upstream_inference_cost`.
650    /// The gateway must not sum both — that would double-count.
651    #[test]
652    fn non_byok_upstream_equal_to_cost_is_not_doubled() {
653        let mut s = Scanner::new(Provider::OpenAi, None);
654        s.feed_body(
655            br#"{"model":"deepseek/deepseek-v4-flash","usage":{"prompt_tokens":50,"completion_tokens":5,"cost":0.000001568,"cost_details":{"upstream_inference_cost":0.000001568}}}"#,
656        );
657        let u = s.finalize().expect("usage");
658        let cost = u.provider_cost_usd.expect("measured cost");
659        assert!(
660            (cost - 0.000001568).abs() < 1e-15,
661            "#746: must book cost once, not 2x; got {cost}"
662        );
663    }
664
665    #[test]
666    fn gateway_header_cost_fills_when_body_has_none() {
667        // #1189: a LiteLLM-style gateway reports the bill in a header; tokens
668        // come from a normal OpenAI-shaped body without `usage.cost`.
669        let mut s = Scanner::new(Provider::OpenAi, None).with_header_cost(Some(0.0042));
670        s.feed_body(
671            br#"{"model":"azure-gpt-4o","usage":{"prompt_tokens":100,"completion_tokens":10}}"#,
672        );
673        let u = s.finalize().expect("usage");
674        assert_eq!(u.provider_cost_usd, Some(0.0042), "header is measured");
675    }
676
677    #[test]
678    fn body_reported_cost_beats_the_header_figure() {
679        // OpenRouter behind another gateway: `usage.cost` is the bill itself.
680        let mut s = Scanner::new(Provider::OpenAi, None).with_header_cost(Some(9.99));
681        s.feed_body(
682            br#"{"model":"deepseek/deepseek-v4-flash","usage":{"prompt_tokens":100,"completion_tokens":10,"cost":0.05}}"#,
683        );
684        let u = s.finalize().expect("usage");
685        assert_eq!(u.provider_cost_usd, Some(0.05), "body wins over header");
686    }
687
688    #[test]
689    fn openrouter_free_model_reports_zero_cost_as_measured() {
690        let mut s = Scanner::new(Provider::OpenAi, None);
691        s.feed_body(
692            br#"{"model":"poolside/laguna-xs-2.1:free","usage":{"prompt_tokens":80,"completion_tokens":20,"cost":0}}"#,
693        );
694        let u = s.finalize().expect("usage");
695        assert_eq!(u.provider_cost_usd, Some(0.0), "free is a price, not a gap");
696    }
697
698    #[test]
699    fn gemini_usage_metadata_with_url_model() {
700        let u = feed_lines(
701            Provider::Gemini,
702            Some("gemini-2.5-pro"),
703            &[
704                r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":25,"candidatesTokenCount":7,"thoughtsTokenCount":39,"totalTokenCount":71}}"#,
705            ],
706        )
707        .expect("usage");
708        assert_eq!(u.model, "gemini-2.5-pro");
709        assert_eq!(u.input_tokens, 25);
710        assert_eq!(u.output_tokens, 46); // candidates 7 + thoughts 39
711        assert_eq!(u.reasoning_tokens, 39);
712    }
713
714    #[test]
715    fn gemini_prefers_model_version_over_url() {
716        let u = feed_lines(
717            Provider::Gemini,
718            Some("gemini-2.5-pro"),
719            &[
720                r#"data: {"modelVersion":"gemini-2.5-pro-002","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"cachedContentTokenCount":4,"totalTokenCount":15}}"#,
721            ],
722        )
723        .expect("usage");
724        assert_eq!(u.model, "gemini-2.5-pro-002");
725        assert_eq!(u.input_tokens, 6); // 10 - 4 cached
726        assert_eq!(u.cache_read_tokens, 4);
727        assert_eq!(u.output_tokens, 5);
728    }
729
730    #[test]
731    fn split_chunks_reassemble_across_feed_calls() {
732        let mut s = Scanner::new(Provider::Anthropic, None);
733        let line = r#"data: {"type":"message_delta","delta":{},"usage":{"output_tokens":42}}"#;
734        let bytes = format!("{line}\n");
735        let (a, b) = bytes.as_bytes().split_at(20);
736        s.feed(a);
737        s.feed(b);
738        let u = s.finalize().expect("usage");
739        assert_eq!(u.output_tokens, 42);
740    }
741
742    #[test]
743    fn no_usage_yields_none() {
744        let out = feed_lines(
745            Provider::OpenAi,
746            None,
747            &[r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"gpt-5.4"}"#],
748        );
749        assert!(out.is_none(), "content-only stream reports no usage");
750    }
751
752    #[test]
753    fn final_event_without_trailing_newline() {
754        let mut s = Scanner::new(Provider::Anthropic, None);
755        s.feed(br#"data: {"type":"message_delta","delta":{},"usage":{"output_tokens":7}}"#);
756        let u = s.finalize().expect("flushes trailing partial line");
757        assert_eq!(u.output_tokens, 7);
758    }
759
760    #[test]
761    fn gemini_model_from_path_extracts() {
762        assert_eq!(
763            gemini_model_from_path("/v1beta/models/gemini-2.5-pro:streamGenerateContent")
764                .as_deref(),
765            Some("gemini-2.5-pro")
766        );
767        assert_eq!(
768            gemini_model_from_path("/v1beta/models/gemini-2.5-flash:generateContent").as_deref(),
769            Some("gemini-2.5-flash")
770        );
771        assert_eq!(gemini_model_from_path("/v1/chat/completions"), None);
772    }
773
774    #[tokio::test]
775    async fn tee_stream_passes_bytes_through_and_records() {
776        let chunks: Vec<Result<Vec<u8>, std::convert::Infallible>> = vec![
777            Ok(b"data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-haiku-4.5\",\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}}\n".to_vec()),
778            Ok(b"data: {\"type\":\"message_delta\",\"delta\":{},\"usage\":{\"output_tokens\":9}}\n".to_vec()),
779        ];
780        let inner = futures::stream::iter(chunks);
781        let scanner = Scanner::new(Provider::Anthropic, None);
782        let teed = tee_stream(inner, scanner);
783        let collected: Vec<_> = teed.collect().await;
784        // Byte-for-byte passthrough preserved.
785        assert_eq!(collected.len(), 2);
786        assert!(collected.iter().all(std::result::Result::is_ok));
787    }
788}