rws_proxy/llm/
usage.rs

1//! Token usage out of provider responses, as they stream past.
2//!
3//! [`UsageParser`] takes a response body in pieces of any size, as a JSON
4//! document or as `text/event-stream`, and keeps the latest usage it saw.
5//! It reads the shapes of OpenAI's Chat Completions (`usage.prompt_tokens`
6//! / `completion_tokens`, in the final chunk of a stream), OpenAI's
7//! Responses API (`usage.input_tokens` / `output_tokens`, nested under
8//! `response` in its `response.completed` event) and Anthropic's Messages
9//! API (`message_start` → `message.usage`, `message_delta` → `usage`, whose
10//! `output_tokens` is cumulative). Anthropic's
11//! `cache_creation_input_tokens` and `cache_read_input_tokens` count as
12//! input: they are billed, so they count toward budgets.
13//!
14//! Loops only, every buffer bounded: provider responses are untrusted input.
15
16use std::io::Read;
17use std::time::{Duration, Instant};
18
19use crate::mcp::json_rpc;
20
21/// Tokens a response reported.
22#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
23pub struct Usage {
24    /// Prompt tokens, cache reads and writes included.
25    pub input_tokens: u64,
26    pub output_tokens: u64,
27}
28
29impl Usage {
30    pub fn total(&self) -> u64 {
31        self.input_tokens.saturating_add(self.output_tokens)
32    }
33}
34
35/// The longest event-stream line kept; a longer one is skipped.
36const MAX_LINE: usize = 1024 * 1024;
37/// The largest JSON body parsed; past it the body is not read for usage.
38const MAX_JSON: usize = 8 * 1024 * 1024;
39
40/// Usage from a response body fed in pieces. See the [module docs](self).
41#[derive(Debug)]
42pub struct UsageParser {
43    sse: bool,
44    buffer: Vec<u8>,
45    /// The current line (SSE) or body (JSON) grew past its cap.
46    overflow: bool,
47    usage: Option<Usage>,
48}
49
50impl UsageParser {
51    /// `content_type` is the response's: `text/event-stream` parses events,
52    /// anything else one JSON document.
53    pub fn new(content_type: &str) -> UsageParser {
54        UsageParser {
55            sse: content_type.to_ascii_lowercase().contains("text/event-stream"),
56            buffer: Vec::new(),
57            overflow: false,
58            usage: None,
59        }
60    }
61
62    pub fn feed(&mut self, mut bytes: &[u8]) {
63        if !self.sse {
64            if !self.overflow && self.buffer.len() + bytes.len() <= MAX_JSON {
65                self.buffer.extend_from_slice(bytes);
66            } else {
67                self.overflow = true;
68                self.buffer = Vec::new();
69            }
70            return;
71        }
72        while let Some(end) = bytes.iter().position(|&b| b == b'\n') {
73            if !self.overflow && self.buffer.len() + end <= MAX_LINE {
74                self.buffer.extend_from_slice(&bytes[..end]);
75                let line = std::mem::take(&mut self.buffer);
76                self.line(&line);
77            }
78            self.buffer.clear();
79            self.overflow = false;
80            bytes = &bytes[end + 1..];
81        }
82        if self.overflow || self.buffer.len() + bytes.len() > MAX_LINE {
83            self.overflow = true;
84            self.buffer.clear();
85        } else {
86            self.buffer.extend_from_slice(bytes);
87        }
88    }
89
90    /// The last usage the body reported, `None` if it reported none.
91    pub fn finish(mut self) -> Option<Usage> {
92        let rest = std::mem::take(&mut self.buffer);
93        if !self.overflow {
94            if self.sse {
95                self.line(&rest);
96            } else if let Ok(text) = std::str::from_utf8(&rest) {
97                self.object(text);
98            }
99        }
100        self.usage
101    }
102
103    fn line(&mut self, line: &[u8]) {
104        let line = line.strip_suffix(b"\r").unwrap_or(line);
105        let Some(data) = line.strip_prefix(b"data:") else { return };
106        let Ok(data) = std::str::from_utf8(data) else { return };
107        let data = data.trim();
108        if data.starts_with('{') {
109            self.object(data);
110        }
111    }
112
113    /// One event or body: its own `usage`, or one under `message` /
114    /// `response`.
115    fn object(&mut self, json: &str) {
116        let Some(members) = json_rpc::object_members(json) else { return };
117        for (key, value) in members {
118            match key.as_str() {
119                "usage" => self.apply(&value),
120                "message" | "response" => {
121                    if let Some(usage) = member(&value, "usage") {
122                        self.apply(&usage);
123                    }
124                }
125                _ => {}
126            }
127        }
128    }
129
130    /// Overwrite the counts a usage object names; keep the others.
131    fn apply(&mut self, usage: &str) {
132        let Some(members) = json_rpc::object_members(usage) else { return };
133        let mut input = None::<u64>;
134        let mut output = None::<u64>;
135        for (key, value) in members {
136            let Ok(n) = value.parse::<u64>() else { continue };
137            match key.as_str() {
138                "prompt_tokens" | "input_tokens" | "cache_creation_input_tokens" | "cache_read_input_tokens" => {
139                    input = Some(input.unwrap_or(0).saturating_add(n));
140                }
141                "completion_tokens" | "output_tokens" => output = Some(n),
142                _ => {}
143            }
144        }
145        if input.is_none() && output.is_none() {
146            return;
147        }
148        let usage = self.usage.get_or_insert_with(Usage::default);
149        if let Some(n) = input {
150            usage.input_tokens = n;
151        }
152        if let Some(n) = output {
153            usage.output_tokens = n;
154        }
155    }
156}
157
158/// The raw value of top-level member `key` of object `json`.
159#[doc(hidden)]
160pub fn member(json: &str, key: &str) -> Option<String> {
161    json_rpc::object_members(json)?.into_iter().find(|(k, _)| k == key).map(|(_, v)| v)
162}
163
164/// HTTP/1.1 chunked framing removed from bytes fed in pieces.
165#[derive(Debug, Default)]
166#[doc(hidden)]
167pub struct ChunkDecoder {
168    state: ChunkState,
169    line: Vec<u8>,
170}
171
172#[derive(Debug, Default, PartialEq)]
173enum ChunkState {
174    #[default]
175    Size,
176    Data(u64),
177    DataEnd,
178    Trailers,
179    Done,
180}
181
182/// The longest chunk-size or trailer line accepted.
183const MAX_CHUNK_LINE: usize = 8 * 1024;
184
185impl ChunkDecoder {
186    #[doc(hidden)]
187    pub fn done(&self) -> bool {
188        self.state == ChunkState::Done
189    }
190
191    /// Append the data in `input` to `out`; an error for broken framing.
192    #[doc(hidden)]
193    pub fn decode(&mut self, mut input: &[u8], out: &mut Vec<u8>) -> Result<(), String> {
194        while !input.is_empty() {
195            match self.state {
196                ChunkState::Data(left) => {
197                    let n = (left.min(input.len() as u64)) as usize;
198                    out.extend_from_slice(&input[..n]);
199                    input = &input[n..];
200                    self.state = if left == n as u64 { ChunkState::DataEnd } else { ChunkState::Data(left - n as u64) };
201                }
202                ChunkState::Done => return Ok(()),
203                ChunkState::Size | ChunkState::DataEnd | ChunkState::Trailers => {
204                    let Some(end) = input.iter().position(|&b| b == b'\n') else {
205                        if self.line.len() + input.len() > MAX_CHUNK_LINE {
206                            return Err("chunk line too long".to_string());
207                        }
208                        self.line.extend_from_slice(input);
209                        return Ok(());
210                    };
211                    if self.line.len() + end > MAX_CHUNK_LINE {
212                        return Err("chunk line too long".to_string());
213                    }
214                    self.line.extend_from_slice(&input[..end]);
215                    input = &input[end + 1..];
216                    let line = std::mem::take(&mut self.line);
217                    let line = line.strip_suffix(b"\r").unwrap_or(&line);
218                    self.state = match self.state {
219                        ChunkState::DataEnd if line.is_empty() => ChunkState::Size,
220                        ChunkState::DataEnd => return Err("chunk data longer than its size".to_string()),
221                        ChunkState::Trailers if line.is_empty() => ChunkState::Done,
222                        ChunkState::Trailers => ChunkState::Trailers,
223                        _ => {
224                            let size = line.split(|&b| b == b';').next().unwrap_or_default();
225                            let size = std::str::from_utf8(size).ok().map(str::trim).filter(|s| !s.is_empty() && s.len() <= 16);
226                            match size.and_then(|s| u64::from_str_radix(s, 16).ok()) {
227                                Some(0) => ChunkState::Trailers,
228                                Some(n) => ChunkState::Data(n),
229                                None => return Err("bad chunk size".to_string()),
230                            }
231                        }
232                    };
233                }
234            }
235        }
236        Ok(())
237    }
238}
239
240/// What a finished response reports: its usage (if any) and the time from
241/// the request to its first body byte.
242#[doc(hidden)]
243pub type OnFinish = Box<dyn FnOnce(Option<Usage>, Option<Duration>) + Send>;
244
245/// A streamed response body read through a [`UsageParser`]: the bytes pass
246/// through (chunked framing removed when `chunked`) and `on_finish` runs
247/// once, at the end of the body or when the reader is dropped early (a
248/// client that went away still spent the tokens it got).
249#[doc(hidden)]
250pub struct UsageReader {
251    inner: Box<dyn Read + Send>,
252    chunked: Option<ChunkDecoder>,
253    parser: Option<UsageParser>,
254    started: Instant,
255    first_byte: Option<Duration>,
256    pending: Vec<u8>,
257    offset: usize,
258    raw: Vec<u8>,
259    on_finish: Option<OnFinish>,
260}
261
262impl UsageReader {
263    #[doc(hidden)]
264    pub fn new(inner: Box<dyn Read + Send>, content_type: &str, chunked: bool, started: Instant, on_finish: OnFinish) -> UsageReader {
265        UsageReader {
266            inner,
267            chunked: chunked.then(ChunkDecoder::default),
268            parser: Some(UsageParser::new(content_type)),
269            started,
270            first_byte: None,
271            pending: Vec::new(),
272            offset: 0,
273            raw: vec![0; 16 * 1024],
274            on_finish: Some(on_finish),
275        }
276    }
277
278    fn finish(&mut self) {
279        if let Some(on_finish) = self.on_finish.take() {
280            let usage = self.parser.take().and_then(UsageParser::finish);
281            on_finish(usage, self.first_byte);
282        }
283    }
284
285    fn saw(first_byte: &mut Option<Duration>, started: Instant, parser: &mut Option<UsageParser>, bytes: &[u8]) {
286        if bytes.is_empty() {
287            return;
288        }
289        first_byte.get_or_insert_with(|| started.elapsed());
290        if let Some(parser) = parser {
291            parser.feed(bytes);
292        }
293    }
294}
295
296impl Read for UsageReader {
297    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
298        if out.is_empty() {
299            return Ok(0);
300        }
301        loop {
302            if self.offset < self.pending.len() {
303                let n = out.len().min(self.pending.len() - self.offset);
304                out[..n].copy_from_slice(&self.pending[self.offset..self.offset + n]);
305                self.offset += n;
306                return Ok(n);
307            }
308            if self.chunked.as_ref().is_some_and(ChunkDecoder::done) {
309                self.finish();
310                return Ok(0);
311            }
312            let Some(decoder) = &mut self.chunked else {
313                let n = self.inner.read(out)?;
314                if n == 0 {
315                    self.finish();
316                } else {
317                    Self::saw(&mut self.first_byte, self.started, &mut self.parser, &out[..n]);
318                }
319                return Ok(n);
320            };
321            let n = self.inner.read(&mut self.raw)?;
322            if n == 0 {
323                self.finish();
324                return Ok(0);
325            }
326            self.pending.clear();
327            self.offset = 0;
328            decoder
329                .decode(&self.raw[..n], &mut self.pending)
330                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
331            Self::saw(&mut self.first_byte, self.started, &mut self.parser, &self.pending);
332        }
333    }
334}
335
336impl Drop for UsageReader {
337    fn drop(&mut self) {
338        self.finish();
339    }
340}
rws_proxy/llm/usage.rs

rws_proxy/llm/
usage.rs
