1use futures::{Stream, StreamExt};
21use serde_json::Value;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum Provider {
26 Anthropic,
27 OpenAi,
30 Gemini,
31}
32
33impl Provider {
34 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#[derive(Debug, Clone, Default, PartialEq)]
49pub struct RealUsage {
50 pub model: String,
51 pub input_tokens: u64,
53 pub output_tokens: u64,
55 pub cache_read_tokens: u64,
57 pub cache_write_tokens: u64,
60 pub reasoning_tokens: u64,
62 pub provider_cost_usd: Option<f64>,
69 pub cohort: Option<super::holdout::Arm>,
73 pub wire: Option<Box<WireContext>>,
78}
79
80#[derive(Debug, Clone, Default, PartialEq)]
86pub struct WireContext {
87 pub provider: String,
89 pub person: Option<String>,
91 pub team: Option<String>,
93 pub project: Option<String>,
95 pub saved_tokens: u64,
98 pub uncompressed_input_tokens: u64,
101 pub is_local: bool,
104 pub routed_from: Option<String>,
107 pub counterfactual: Option<super::counterfactual::CounterfactualSlot>,
112}
113
114impl RealUsage {
115 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
128const MAX_LINE_BYTES: usize = 1 << 20; pub struct Scanner {
138 provider: Provider,
139 url_model: Option<String>,
141 cohort: Option<super::holdout::Arm>,
143 wire: Option<Box<WireContext>>,
145 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 #[must_use]
167 pub fn with_cohort(mut self, cohort: Option<super::holdout::Arm>) -> Self {
168 self.cohort = cohort;
169 self
170 }
171
172 #[must_use]
175 pub fn with_wire_context(mut self, wire: Option<Box<WireContext>>) -> Self {
176 self.wire = wire;
177 self
178 }
179
180 #[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 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(); 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 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 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 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 if !self.line_might_be_relevant(trimmed) {
246 return;
247 }
248 let json_str = if let Some(rest) = trimmed.strip_prefix("data:") {
249 let r = rest.trim();
251 if r.is_empty() || r == "[DONE]" {
252 return;
253 }
254 r
255 } else if trimmed.starts_with('{') {
256 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 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
288fn 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
321fn 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 return;
346 }
347
348 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 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 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
412fn 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
455pub 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
467pub 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 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); 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 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 #[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 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 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); 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); 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 assert_eq!(collected.len(), 2);
786 assert!(collected.iter().all(std::result::Result::is_ok));
787 }
788}