From 3b895895892f6db5ad0dc11a62adfba2be65a236 Mon Sep 17 00:00:00 2001 From: Mark Stuart Date: Fri, 25 Sep 2026 22:36:04 +0000 Subject: [PATCH] fix(proxy): parse usage from protocol fields --- crates/tare-proxy/src/server.rs | 128 ++++++++++++++++++++++++-------- 1 file changed, 95 insertions(+), 33 deletions(-) diff --git a/crates/tare-proxy/src/server.rs b/crates/tare-proxy/src/server.rs index 12d32a6..c952170 100644 --- a/crates/tare-proxy/src/server.rs +++ b/crates/tare-proxy/src/server.rs @@ -129,50 +129,83 @@ fn detect_anthropic_provider(req: &serde_json::Value) -> Provider { } } -fn scan_u64(s: &str, key: &str) -> Option { - let i = s.find(key)?; - let rest = &s[i + key.len()..]; - let colon = rest.find(':')?; - let after = &rest[colon + 1..]; - let digits: String = after - .chars() - .skip_while(|c| c.is_whitespace()) - .take_while(|c| c.is_ascii_digit()) - .collect(); - digits.parse().ok() +fn for_each_response_json(buf: &[u8], mut visit: impl FnMut(&serde_json::Value)) { + if let Ok(value) = serde_json::from_slice::(buf) { + visit(&value); + return; + } + + for line in buf.split(|byte| *byte == b'\n') { + let line = line.strip_suffix(b"\r").unwrap_or(line); + let payload = line.strip_prefix(b"data:").unwrap_or(line); + let payload = payload + .iter() + .position(|byte| !byte.is_ascii_whitespace()) + .map(|start| &payload[start..]) + .unwrap_or_default(); + if payload.is_empty() || payload == b"[DONE]" { + continue; + } + if let Ok(value) = serde_json::from_slice::(payload) { + visit(&value); + } + } } -/// Parse Anthropic cache usage from a (capped) response buffer — works for both the non-streaming -/// top-level `usage` and the streaming `message_start` event (both carry these keys, near the start -/// of the stream, so the head-capped buffer always contains them). -fn parse_anthropic_usage(buf: &[u8]) -> Option<(u64, u64)> { - let s = String::from_utf8_lossy(buf); - Some(( - scan_u64(&s, "\"cache_read_input_tokens\"")?, - scan_u64(&s, "\"cache_creation_input_tokens\"")?, - )) +fn response_usage( + value: &serde_json::Value, +) -> Option<&serde_json::Map> { + value + .get("usage") + .or_else(|| { + value + .get("message") + .and_then(|message| message.get("usage")) + })? + .as_object() } -/// Largest integer value across ALL occurrences of `key` in `s`. Output token counts are reported -/// cumulatively across streaming events (the final event carries the total), so the max occurrence -/// is the turn total. -fn scan_u64_max(s: &str, key: &str) -> Option { - s.match_indices(key) - .filter_map(|(i, _)| scan_u64(&s[i..], key)) - .max() +/// Parse Anthropic cache usage from a (capped) response buffer — works for both the non-streaming +/// top-level `usage` and the streaming `message_start.message.usage` object. Restricting extraction +/// to those protocol fields prevents assistant text or tool payloads from spoofing usage counters. +fn parse_anthropic_usage(buf: &[u8]) -> Option<(u64, u64)> { + let mut result = None; + for_each_response_json(buf, |value| { + let Some(usage) = response_usage(value) else { + return; + }; + let read = usage + .get("cache_read_input_tokens") + .and_then(|v| v.as_u64()); + let creation = usage + .get("cache_creation_input_tokens") + .and_then(|v| v.as_u64()); + if let (Some(read), Some(creation)) = (read, creation) { + result = Some((read, creation)); + } + }); + result } /// Parse the turn's OUTPUT token count from a (capped) response buffer — `output_tokens` (Anthropic) /// or `completion_tokens` (OpenAI). The total lands in the stream's FINAL usage event; for responses /// under the 2 MB scan cap (essentially all single turns) that event is in the buffer. fn parse_output_tokens(buf: &[u8], provider: Provider) -> Option { - let s = String::from_utf8_lossy(buf); let key = if matches!(provider, Provider::OpenAi) { - "\"completion_tokens\"" + "completion_tokens" } else { - "\"output_tokens\"" + "output_tokens" }; - scan_u64_max(&s, key) + let mut result = None; + for_each_response_json(buf, |value| { + let output_tokens = response_usage(value) + .and_then(|usage| usage.get(key)) + .and_then(|value| value.as_u64()); + if let Some(output_tokens) = output_tokens { + result = Some(result.unwrap_or(0).max(output_tokens)); + } + }); + result } fn hit_rate(read: u64, creation: u64) -> Option { @@ -648,9 +681,10 @@ mod tests { use axum::http::{HeaderMap, HeaderName, HeaderValue}; use super::{ - record_usage_chunk, response_connection_tokens, should_forward_response_header, - TAIL_SCAN_CAP, USAGE_SCAN_CAP, + parse_anthropic_usage, parse_output_tokens, record_usage_chunk, response_connection_tokens, + should_forward_response_header, TAIL_SCAN_CAP, USAGE_SCAN_CAP, }; + use tare_cache::Provider; /// Pins the 413 discrimination used in handle_messages / handle_chat: when axum::body::to_bytes /// hits MAX_BODY_BYTES, the error's Display must contain "length limit". If a future bump to @@ -690,6 +724,34 @@ mod tests { assert!(tail.make_contiguous().ends_with(b"final")); } + #[test] + fn usage_parsing_ignores_field_names_in_model_output() { + let anthropic = br#"{"content":[{"type":"text","text":"\"cache_read_input_tokens\":999999, \"cache_creation_input_tokens\":0, \"output_tokens\":999999"}],"usage":{"cache_read_input_tokens":40,"cache_creation_input_tokens":10,"output_tokens":25}}"#; + assert_eq!(parse_anthropic_usage(anthropic), Some((40, 10))); + assert_eq!( + parse_output_tokens(anthropic, Provider::Anthropic5m), + Some(25) + ); + + let openai = br#"{"choices":[{"message":{"content":"\"completion_tokens\":999999"}}],"usage":{"completion_tokens":30}}"#; + assert_eq!(parse_output_tokens(openai, Provider::OpenAi), Some(30)); + } + + #[test] + fn usage_parsing_reads_sse_protocol_objects_only() { + let anthropic = b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"\\\"output_tokens\\\":999999\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":42}}\n\n"; + assert_eq!( + parse_output_tokens(anthropic, Provider::Anthropic5m), + Some(42) + ); + + let message_start = b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"cache_read_input_tokens\":80,\"cache_creation_input_tokens\":20,\"output_tokens\":1}}}\n\n"; + assert_eq!(parse_anthropic_usage(message_start), Some((80, 20))); + + let openai = b"data: {\"choices\":[{\"delta\":{\"content\":\"\\\"completion_tokens\\\":999999\"}}]}\n\ndata: {\"choices\":[],\"usage\":{\"completion_tokens\":55}}\n\ndata: [DONE]\n\n"; + assert_eq!(parse_output_tokens(openai, Provider::OpenAi), Some(55)); + } + #[test] fn response_header_filter_removes_standard_and_connection_named_hop_headers() { let mut headers = HeaderMap::new();