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    /// Request-side gateway context (enterprise#11/#17/#18): identity tags,
66    /// compression savings and baseline inputs, stamped from the request before
67    /// it left for the upstream. `None` outside the forward path (e.g. tests
68    /// that only parse response bodies).
69    pub wire: Option<Box<WireContext>>,
70}
71
72/// Request-side context the forward path knows and the response scanner does
73/// not: who sent the request (gateway identity, enterprise#11), what the proxy
74/// saved on the wire, and the counterfactual-baseline inputs (enterprise#18).
75/// Travels inside [`RealUsage`] so `usage_meter::record` stays the single
76/// choke-point through which every measured turn flows.
77#[derive(Debug, Clone, Default, PartialEq)]
78pub struct WireContext {
79    /// Provider label (`Anthropic|OpenAI|ChatGPT|Gemini`) of the serving route.
80    pub provider: String,
81    /// Person tag from the gateway key (enterprise#11).
82    pub person: Option<String>,
83    /// Team tag from the gateway key.
84    pub team: Option<String>,
85    /// Project: `x-leanctx-project` header wins over the key's default project.
86    pub project: Option<String>,
87    /// Estimated tokens this request saved through wire compression
88    /// (bytes/4 heuristic, same basis as the proxy stats).
89    pub saved_tokens: u64,
90    /// Estimated request tokens BEFORE lean-ctx compression (bytes/4) — the
91    /// SEE-attribution input of the avoided-cost baseline (enterprise#18).
92    pub uncompressed_input_tokens: u64,
93    /// True when the serving upstream is a local/loopback endpoint — billed via
94    /// the transparent `local_shadow_rate`, never 0 (enterprise#15/#18).
95    pub is_local: bool,
96    /// Originally requested model when the router downgraded/aliased it
97    /// (enterprise#13); `None` for passthrough.
98    pub routed_from: Option<String>,
99    /// Provider-counted input tokens of the ORIGINAL (uncompressed) request,
100    /// from the free `count_tokens` probe (#701). `None` when metering is off
101    /// or no probe was spawned; an empty slot at read time (probe failed or
102    /// still in flight) degrades the row to the local estimate.
103    pub counterfactual: Option<super::counterfactual::CounterfactualSlot>,
104}
105
106impl RealUsage {
107    /// True once any model or token field has been observed — the gate for
108    /// recording. Avoids emitting empty rows for streams that never reported usage.
109    fn is_meaningful(&self) -> bool {
110        !self.model.is_empty()
111            || self.input_tokens > 0
112            || self.output_tokens > 0
113            || self.cache_read_tokens > 0
114            || self.cache_write_tokens > 0
115    }
116}
117
118/// Upper bound on a single buffered line before we give up on it. Usage events
119/// are tiny; this only guards against a pathological newline-free stream.
120const MAX_LINE_BYTES: usize = 1 << 20; // 1 MiB
121
122/// Incrementally extracts [`RealUsage`] from a response stream (or a full body).
123///
124/// `feed` is called with raw response chunks and keeps only the trailing partial
125/// line buffered (O(1) memory beyond one line); `finalize` returns the merged
126/// usage once the stream ends.
127pub struct Scanner {
128    provider: Provider,
129    /// Model parsed from the request URL (Gemini puts it there, not in the body).
130    url_model: Option<String>,
131    /// Output-savings arm (#895), stamped onto the usage at finalize.
132    cohort: Option<super::holdout::Arm>,
133    /// Request-side gateway context (enterprise#11/#18), stamped at finalize.
134    wire: Option<Box<WireContext>>,
135    buf: Vec<u8>,
136    usage: RealUsage,
137}
138
139impl Scanner {
140    pub fn new(provider: Provider, url_model: Option<String>) -> Self {
141        Self {
142            provider,
143            url_model,
144            cohort: None,
145            wire: None,
146            buf: Vec::new(),
147            usage: RealUsage::default(),
148        }
149    }
150
151    /// Tags the usage this scanner produces with an output-savings arm (#895).
152    #[must_use]
153    pub fn with_cohort(mut self, cohort: Option<super::holdout::Arm>) -> Self {
154        self.cohort = cohort;
155        self
156    }
157
158    /// Attaches the request-side gateway context (identity tags, wire savings,
159    /// baseline inputs — enterprise#11/#17/#18) stamped onto the usage record.
160    #[must_use]
161    pub fn with_wire_context(mut self, wire: Option<Box<WireContext>>) -> Self {
162        self.wire = wire;
163        self
164    }
165
166    /// Feeds a raw streaming chunk, scanning every complete line it completes.
167    pub fn feed(&mut self, chunk: &[u8]) {
168        self.buf.extend_from_slice(chunk);
169        while let Some(nl) = self.buf.iter().position(|&b| b == b'\n') {
170            let mut line: Vec<u8> = self.buf.drain(..=nl).collect();
171            line.pop(); // drop '\n'
172            if line.last() == Some(&b'\r') {
173                line.pop();
174            }
175            self.scan_line(&line);
176        }
177        if self.buf.len() > MAX_LINE_BYTES {
178            self.buf.clear();
179        }
180    }
181
182    /// Feeds a complete non-streaming JSON response body.
183    pub fn feed_body(&mut self, body: &[u8]) {
184        if let Ok(v) = serde_json::from_slice::<Value>(body) {
185            self.absorb(&v);
186        }
187    }
188
189    /// Consumes the scanner, flushing any trailing partial line (a final event
190    /// may arrive without a newline) and returning the merged usage if any.
191    pub fn finalize(mut self) -> Option<RealUsage> {
192        if !self.buf.is_empty() {
193            let line = std::mem::take(&mut self.buf);
194            self.scan_line(&line);
195        }
196        if self.usage.is_meaningful() {
197            self.usage.cohort = self.cohort;
198            self.usage.wire = self.wire;
199            Some(self.usage)
200        } else {
201            None
202        }
203    }
204
205    fn scan_line(&mut self, line: &[u8]) {
206        let Ok(text) = std::str::from_utf8(line) else {
207            return;
208        };
209        let trimmed = text.trim();
210        if trimmed.is_empty() {
211            return;
212        }
213        // Cheap pre-filter: skip the bulk of the stream (content deltas) and only
214        // JSON-parse lines that can carry usage or the model name.
215        if !self.line_might_be_relevant(trimmed) {
216            return;
217        }
218        let json_str = if let Some(rest) = trimmed.strip_prefix("data:") {
219            // SSE (Anthropic, OpenAI, Gemini with alt=sse).
220            let r = rest.trim();
221            if r.is_empty() || r == "[DONE]" {
222                return;
223            }
224            r
225        } else if trimmed.starts_with('{') {
226            // NDJSON / array-element line (Gemini x-ndjson). Tolerate the array
227            // punctuation a streamed JSON array puts around an element.
228            trimmed
229                .trim_start_matches([',', '['])
230                .trim_end_matches([',', ']'])
231                .trim()
232        } else {
233            return;
234        };
235        if let Ok(v) = serde_json::from_str::<Value>(json_str) {
236            self.absorb(&v);
237        }
238    }
239
240    fn line_might_be_relevant(&self, s: &str) -> bool {
241        match self.provider {
242            // Anthropic `message_start`/`message_delta` and OpenAI `usage`/
243            // `response.*` events all contain the substring "usage".
244            Provider::Anthropic | Provider::OpenAi => s.contains("usage"),
245            Provider::Gemini => s.contains("usageMetadata"),
246        }
247    }
248
249    fn absorb(&mut self, v: &Value) {
250        match self.provider {
251            Provider::Anthropic => absorb_anthropic(&mut self.usage, v),
252            Provider::OpenAi => absorb_openai(&mut self.usage, v),
253            Provider::Gemini => absorb_gemini(&mut self.usage, v, self.url_model.as_deref()),
254        }
255    }
256}
257
258/// Anthropic: model + input/cache live on `message` (streaming `message_start`
259/// or a non-streaming body); `output_tokens` arrives later on the event-level
260/// `usage` of `message_delta`. Latest non-zero wins, so the cumulative final
261/// delta is authoritative.
262fn absorb_anthropic(u: &mut RealUsage, v: &Value) {
263    let msg = v.get("message").unwrap_or(v);
264    if let Some(model) = msg.get("model").and_then(Value::as_str)
265        && !model.is_empty()
266    {
267        u.model = model.to_string();
268    }
269    let Some(usage) = msg.get("usage").or_else(|| v.get("usage")) else {
270        return;
271    };
272    if let Some(n) = usage.get("input_tokens").and_then(Value::as_u64) {
273        u.input_tokens = n;
274    }
275    if let Some(n) = usage.get("cache_read_input_tokens").and_then(Value::as_u64) {
276        u.cache_read_tokens = n;
277    }
278    if let Some(n) = usage
279        .get("cache_creation_input_tokens")
280        .and_then(Value::as_u64)
281    {
282        u.cache_write_tokens = n;
283    }
284    if let Some(n) = usage.get("output_tokens").and_then(Value::as_u64)
285        && n > 0
286    {
287        u.output_tokens = n;
288    }
289}
290
291/// OpenAI Chat Completions + Responses. `response.completed` nests the payload
292/// under `response`; chat chunks and non-streaming bodies are top-level. Both
293/// `usage` dialects are accepted (Responses: `input_tokens`/`output_tokens`;
294/// Chat: `prompt_tokens`/`completion_tokens`). `cached_tokens` is the cache-read
295/// portion of the reported input; OpenAI bills no separate cache write.
296fn absorb_openai(u: &mut RealUsage, v: &Value) {
297    let root = v.get("response").unwrap_or(v);
298    if let Some(model) = root.get("model").and_then(Value::as_str)
299        && !model.is_empty()
300    {
301        u.model = model.to_string();
302    }
303    let Some(usage) = root.get("usage") else {
304        return;
305    };
306    if usage.is_null() {
307        // `response.created` / `response.in_progress` carry `usage: null`.
308        return;
309    }
310
311    let total_input = usage
312        .get("input_tokens")
313        .or_else(|| usage.get("prompt_tokens"))
314        .and_then(Value::as_u64)
315        .unwrap_or(0);
316    let total_output = usage
317        .get("output_tokens")
318        .or_else(|| usage.get("completion_tokens"))
319        .and_then(Value::as_u64)
320        .unwrap_or(0);
321    let cached = usage
322        .get("input_tokens_details")
323        .or_else(|| usage.get("prompt_tokens_details"))
324        .and_then(|d| d.get("cached_tokens"))
325        .and_then(Value::as_u64)
326        .unwrap_or(0);
327    let reasoning = usage
328        .get("output_tokens_details")
329        .or_else(|| usage.get("completion_tokens_details"))
330        .and_then(|d| d.get("reasoning_tokens"))
331        .and_then(Value::as_u64)
332        .unwrap_or(0);
333
334    if total_input == 0 && total_output == 0 {
335        return;
336    }
337    u.input_tokens = total_input.saturating_sub(cached);
338    u.cache_read_tokens = cached;
339    u.cache_write_tokens = 0;
340    u.output_tokens = total_output;
341    u.reasoning_tokens = reasoning;
342}
343
344/// Gemini: `usageMetadata` carries the counts; `modelVersion` (or the request
345/// URL) the model. `thoughtsTokenCount` is billed at the output rate and is not
346/// part of `candidatesTokenCount`, so it is added into `output_tokens`.
347fn absorb_gemini(u: &mut RealUsage, v: &Value, url_model: Option<&str>) {
348    if let Some(mv) = v.get("modelVersion").and_then(Value::as_str)
349        && !mv.is_empty()
350    {
351        u.model = mv.to_string();
352    } else if u.model.is_empty()
353        && let Some(m) = url_model
354        && !m.is_empty()
355    {
356        u.model = m.to_string();
357    }
358    let Some(um) = v.get("usageMetadata") else {
359        return;
360    };
361    let prompt = um
362        .get("promptTokenCount")
363        .and_then(Value::as_u64)
364        .unwrap_or(0);
365    let candidates = um
366        .get("candidatesTokenCount")
367        .and_then(Value::as_u64)
368        .unwrap_or(0);
369    let cached = um
370        .get("cachedContentTokenCount")
371        .and_then(Value::as_u64)
372        .unwrap_or(0);
373    let thoughts = um
374        .get("thoughtsTokenCount")
375        .and_then(Value::as_u64)
376        .unwrap_or(0);
377    if prompt == 0 && candidates == 0 && thoughts == 0 {
378        return;
379    }
380    u.input_tokens = prompt.saturating_sub(cached);
381    u.cache_read_tokens = cached;
382    u.cache_write_tokens = 0;
383    u.output_tokens = candidates + thoughts;
384    u.reasoning_tokens = thoughts;
385}
386
387/// Extracts the model from a Gemini request path
388/// (`/v1beta/models/{model}:generateContent`). Returns `None` for other paths.
389pub fn gemini_model_from_path(path: &str) -> Option<String> {
390    let after = path.rsplit_once("/models/").map(|(_, m)| m)?;
391    let model = after.split(':').next().unwrap_or(after).trim();
392    if model.is_empty() {
393        None
394    } else {
395        Some(model.to_string())
396    }
397}
398
399/// Wraps a response byte stream so every chunk is forwarded **byte-for-byte**
400/// while a [`Scanner`] observes it; on stream end the merged usage is recorded.
401/// Memory overhead is one buffered SSE line, never the whole response.
402pub fn tee_stream<S, B, E>(
403    inner: S,
404    scanner: Scanner,
405) -> impl Stream<Item = Result<B, E>> + Send + 'static
406where
407    S: Stream<Item = Result<B, E>> + Send + Unpin + 'static,
408    B: AsRef<[u8]> + Send + 'static,
409    E: Send + 'static,
410{
411    futures::stream::unfold(
412        (inner, Some(scanner)),
413        |(mut inner, mut scanner)| async move {
414            match inner.next().await {
415                Some(Ok(chunk)) => {
416                    if let Some(s) = scanner.as_mut() {
417                        s.feed(chunk.as_ref());
418                    }
419                    Some((Ok(chunk), (inner, scanner)))
420                }
421                Some(err) => Some((err, (inner, scanner))),
422                None => {
423                    if let Some(s) = scanner.take()
424                        && let Some(usage) = s.finalize()
425                    {
426                        super::usage_meter::record(&usage);
427                    }
428                    None
429                }
430            }
431        },
432    )
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438
439    fn feed_lines(
440        provider: Provider,
441        url_model: Option<&str>,
442        lines: &[&str],
443    ) -> Option<RealUsage> {
444        let mut s = Scanner::new(provider, url_model.map(str::to_string));
445        for line in lines {
446            s.feed(line.as_bytes());
447            s.feed(b"\n");
448        }
449        s.finalize()
450    }
451
452    #[test]
453    fn anthropic_merges_message_start_and_delta() {
454        let u = feed_lines(
455            Provider::Anthropic,
456            None,
457            &[
458                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}}}"#,
459                r#"data: {"type":"content_block_delta","index":0,"delta":{"text":"hello"}}"#,
460                r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":73}}"#,
461                "data: {\"type\":\"message_stop\"}",
462            ],
463        )
464        .expect("usage");
465        assert_eq!(u.model, "claude-opus-4-5-20251101");
466        assert_eq!(u.input_tokens, 100);
467        assert_eq!(u.cache_read_tokens, 2000);
468        assert_eq!(u.cache_write_tokens, 50);
469        assert_eq!(u.output_tokens, 73);
470    }
471
472    #[test]
473    fn anthropic_non_streaming_body() {
474        let mut s = Scanner::new(Provider::Anthropic, None);
475        s.feed_body(
476            br#"{"model":"claude-sonnet-4-5","usage":{"input_tokens":24,"output_tokens":18,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}"#,
477        );
478        let u = s.finalize().expect("usage");
479        assert_eq!(u.model, "claude-sonnet-4-5");
480        assert_eq!(u.input_tokens, 24);
481        assert_eq!(u.output_tokens, 18);
482    }
483
484    #[test]
485    fn scanner_stamps_wire_context_onto_usage() {
486        // enterprise#11/#17: the request-side context must survive scanning and
487        // arrive on the finalized record that usage_meter/store consume.
488        let wire = Box::new(WireContext {
489            provider: "Anthropic".into(),
490            person: Some("yves".into()),
491            team: None,
492            project: Some("billing".into()),
493            saved_tokens: 42,
494            uncompressed_input_tokens: 500,
495            is_local: false,
496            routed_from: None,
497            counterfactual: None,
498        });
499        let mut s = Scanner::new(Provider::Anthropic, None).with_wire_context(Some(wire.clone()));
500        s.feed_body(
501            br#"{"model":"claude-sonnet-4-5","usage":{"input_tokens":24,"output_tokens":18}}"#,
502        );
503        let u = s.finalize().expect("usage");
504        assert_eq!(u.wire, Some(wire));
505    }
506
507    #[test]
508    fn openai_responses_completed_event() {
509        let u = feed_lines(
510            Provider::OpenAi,
511            None,
512            &[
513                r#"data: {"type":"response.created","response":{"model":"gpt-5.4","usage":null}}"#,
514                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}}}"#,
515            ],
516        )
517        .expect("usage");
518        assert_eq!(u.model, "gpt-5.4");
519        assert_eq!(u.input_tokens, 1000); // 1289 - 289 cached
520        assert_eq!(u.cache_read_tokens, 289);
521        assert_eq!(u.output_tokens, 685);
522        assert_eq!(u.reasoning_tokens, 640);
523    }
524
525    #[test]
526    fn openai_chat_final_usage_chunk() {
527        let u = feed_lines(
528            Provider::OpenAi,
529            None,
530            &[
531                r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"gpt-5.4-mini"}"#,
532                r#"data: {"choices":[],"model":"gpt-5.4-mini","usage":{"prompt_tokens":500,"prompt_tokens_details":{"cached_tokens":100},"completion_tokens":40,"total_tokens":540}}"#,
533                "data: [DONE]",
534            ],
535        )
536        .expect("usage");
537        assert_eq!(u.model, "gpt-5.4-mini");
538        assert_eq!(u.input_tokens, 400);
539        assert_eq!(u.cache_read_tokens, 100);
540        assert_eq!(u.output_tokens, 40);
541    }
542
543    #[test]
544    fn gemini_usage_metadata_with_url_model() {
545        let u = feed_lines(
546            Provider::Gemini,
547            Some("gemini-2.5-pro"),
548            &[
549                r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":25,"candidatesTokenCount":7,"thoughtsTokenCount":39,"totalTokenCount":71}}"#,
550            ],
551        )
552        .expect("usage");
553        assert_eq!(u.model, "gemini-2.5-pro");
554        assert_eq!(u.input_tokens, 25);
555        assert_eq!(u.output_tokens, 46); // candidates 7 + thoughts 39
556        assert_eq!(u.reasoning_tokens, 39);
557    }
558
559    #[test]
560    fn gemini_prefers_model_version_over_url() {
561        let u = feed_lines(
562            Provider::Gemini,
563            Some("gemini-2.5-pro"),
564            &[
565                r#"data: {"modelVersion":"gemini-2.5-pro-002","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"cachedContentTokenCount":4,"totalTokenCount":15}}"#,
566            ],
567        )
568        .expect("usage");
569        assert_eq!(u.model, "gemini-2.5-pro-002");
570        assert_eq!(u.input_tokens, 6); // 10 - 4 cached
571        assert_eq!(u.cache_read_tokens, 4);
572        assert_eq!(u.output_tokens, 5);
573    }
574
575    #[test]
576    fn split_chunks_reassemble_across_feed_calls() {
577        let mut s = Scanner::new(Provider::Anthropic, None);
578        let line = r#"data: {"type":"message_delta","delta":{},"usage":{"output_tokens":42}}"#;
579        let bytes = format!("{line}\n");
580        let (a, b) = bytes.as_bytes().split_at(20);
581        s.feed(a);
582        s.feed(b);
583        let u = s.finalize().expect("usage");
584        assert_eq!(u.output_tokens, 42);
585    }
586
587    #[test]
588    fn no_usage_yields_none() {
589        let out = feed_lines(
590            Provider::OpenAi,
591            None,
592            &[r#"data: {"choices":[{"delta":{"content":"hi"}}],"model":"gpt-5.4"}"#],
593        );
594        assert!(out.is_none(), "content-only stream reports no usage");
595    }
596
597    #[test]
598    fn final_event_without_trailing_newline() {
599        let mut s = Scanner::new(Provider::Anthropic, None);
600        s.feed(br#"data: {"type":"message_delta","delta":{},"usage":{"output_tokens":7}}"#);
601        let u = s.finalize().expect("flushes trailing partial line");
602        assert_eq!(u.output_tokens, 7);
603    }
604
605    #[test]
606    fn gemini_model_from_path_extracts() {
607        assert_eq!(
608            gemini_model_from_path("/v1beta/models/gemini-2.5-pro:streamGenerateContent")
609                .as_deref(),
610            Some("gemini-2.5-pro")
611        );
612        assert_eq!(
613            gemini_model_from_path("/v1beta/models/gemini-2.5-flash:generateContent").as_deref(),
614            Some("gemini-2.5-flash")
615        );
616        assert_eq!(gemini_model_from_path("/v1/chat/completions"), None);
617    }
618
619    #[tokio::test]
620    async fn tee_stream_passes_bytes_through_and_records() {
621        let chunks: Vec<Result<Vec<u8>, std::convert::Infallible>> = vec![
622            Ok(b"data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-haiku-4.5\",\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}}\n".to_vec()),
623            Ok(b"data: {\"type\":\"message_delta\",\"delta\":{},\"usage\":{\"output_tokens\":9}}\n".to_vec()),
624        ];
625        let inner = futures::stream::iter(chunks);
626        let scanner = Scanner::new(Provider::Anthropic, None);
627        let teed = tee_stream(inner, scanner);
628        let collected: Vec<_> = teed.collect().await;
629        // Byte-for-byte passthrough preserved.
630        assert_eq!(collected.len(), 2);
631        assert!(collected.iter().all(std::result::Result::is_ok));
632    }
633}
lean_ctx/proxy/usage.rs

lean_ctx/proxy/
usage.rs
