rust_web_server/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`.
159pub(crate) fn member(json: &str, key: &str) -> Option<String> {
160    json_rpc::object_members(json)?.into_iter().find(|(k, _)| k == key).map(|(_, v)| v)
161}
162
163/// HTTP/1.1 chunked framing removed from bytes fed in pieces.
164#[derive(Debug, Default)]
165pub(crate) struct ChunkDecoder {
166    state: ChunkState,
167    line: Vec<u8>,
168}
169
170#[derive(Debug, Default, PartialEq)]
171enum ChunkState {
172    #[default]
173    Size,
174    Data(u64),
175    DataEnd,
176    Trailers,
177    Done,
178}
179
180/// The longest chunk-size or trailer line accepted.
181const MAX_CHUNK_LINE: usize = 8 * 1024;
182
183impl ChunkDecoder {
184    pub(crate) fn done(&self) -> bool {
185        self.state == ChunkState::Done
186    }
187
188    /// Append the data in `input` to `out`; an error for broken framing.
189    pub(crate) fn decode(&mut self, mut input: &[u8], out: &mut Vec<u8>) -> Result<(), String> {
190        while !input.is_empty() {
191            match self.state {
192                ChunkState::Data(left) => {
193                    let n = (left.min(input.len() as u64)) as usize;
194                    out.extend_from_slice(&input[..n]);
195                    input = &input[n..];
196                    self.state = if left == n as u64 { ChunkState::DataEnd } else { ChunkState::Data(left - n as u64) };
197                }
198                ChunkState::Done => return Ok(()),
199                ChunkState::Size | ChunkState::DataEnd | ChunkState::Trailers => {
200                    let Some(end) = input.iter().position(|&b| b == b'\n') else {
201                        if self.line.len() + input.len() > MAX_CHUNK_LINE {
202                            return Err("chunk line too long".to_string());
203                        }
204                        self.line.extend_from_slice(input);
205                        return Ok(());
206                    };
207                    if self.line.len() + end > MAX_CHUNK_LINE {
208                        return Err("chunk line too long".to_string());
209                    }
210                    self.line.extend_from_slice(&input[..end]);
211                    input = &input[end + 1..];
212                    let line = std::mem::take(&mut self.line);
213                    let line = line.strip_suffix(b"\r").unwrap_or(&line);
214                    self.state = match self.state {
215                        ChunkState::DataEnd if line.is_empty() => ChunkState::Size,
216                        ChunkState::DataEnd => return Err("chunk data longer than its size".to_string()),
217                        ChunkState::Trailers if line.is_empty() => ChunkState::Done,
218                        ChunkState::Trailers => ChunkState::Trailers,
219                        _ => {
220                            let size = line.split(|&b| b == b';').next().unwrap_or_default();
221                            let size = std::str::from_utf8(size).ok().map(str::trim).filter(|s| !s.is_empty() && s.len() <= 16);
222                            match size.and_then(|s| u64::from_str_radix(s, 16).ok()) {
223                                Some(0) => ChunkState::Trailers,
224                                Some(n) => ChunkState::Data(n),
225                                None => return Err("bad chunk size".to_string()),
226                            }
227                        }
228                    };
229                }
230            }
231        }
232        Ok(())
233    }
234}
235
236/// What a finished response reports: its usage (if any) and the time from
237/// the request to its first body byte.
238pub(crate) type OnFinish = Box<dyn FnOnce(Option<Usage>, Option<Duration>) + Send>;
239
240/// A streamed response body read through a [`UsageParser`]: the bytes pass
241/// through (chunked framing removed when `chunked`) and `on_finish` runs
242/// once, at the end of the body or when the reader is dropped early (a
243/// client that went away still spent the tokens it got).
244pub(crate) struct UsageReader {
245    inner: Box<dyn Read + Send>,
246    chunked: Option<ChunkDecoder>,
247    parser: Option<UsageParser>,
248    started: Instant,
249    first_byte: Option<Duration>,
250    pending: Vec<u8>,
251    offset: usize,
252    raw: Vec<u8>,
253    on_finish: Option<OnFinish>,
254}
255
256impl UsageReader {
257    pub(crate) fn new(inner: Box<dyn Read + Send>, content_type: &str, chunked: bool, started: Instant, on_finish: OnFinish) -> UsageReader {
258        UsageReader {
259            inner,
260            chunked: chunked.then(ChunkDecoder::default),
261            parser: Some(UsageParser::new(content_type)),
262            started,
263            first_byte: None,
264            pending: Vec::new(),
265            offset: 0,
266            raw: vec![0; 16 * 1024],
267            on_finish: Some(on_finish),
268        }
269    }
270
271    fn finish(&mut self) {
272        if let Some(on_finish) = self.on_finish.take() {
273            let usage = self.parser.take().and_then(UsageParser::finish);
274            on_finish(usage, self.first_byte);
275        }
276    }
277
278    fn saw(first_byte: &mut Option<Duration>, started: Instant, parser: &mut Option<UsageParser>, bytes: &[u8]) {
279        if bytes.is_empty() {
280            return;
281        }
282        first_byte.get_or_insert_with(|| started.elapsed());
283        if let Some(parser) = parser {
284            parser.feed(bytes);
285        }
286    }
287}
288
289impl Read for UsageReader {
290    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
291        if out.is_empty() {
292            return Ok(0);
293        }
294        loop {
295            if self.offset < self.pending.len() {
296                let n = out.len().min(self.pending.len() - self.offset);
297                out[..n].copy_from_slice(&self.pending[self.offset..self.offset + n]);
298                self.offset += n;
299                return Ok(n);
300            }
301            if self.chunked.as_ref().is_some_and(ChunkDecoder::done) {
302                self.finish();
303                return Ok(0);
304            }
305            let Some(decoder) = &mut self.chunked else {
306                let n = self.inner.read(out)?;
307                if n == 0 {
308                    self.finish();
309                } else {
310                    Self::saw(&mut self.first_byte, self.started, &mut self.parser, &out[..n]);
311                }
312                return Ok(n);
313            };
314            let n = self.inner.read(&mut self.raw)?;
315            if n == 0 {
316                self.finish();
317                return Ok(0);
318            }
319            self.pending.clear();
320            self.offset = 0;
321            decoder
322                .decode(&self.raw[..n], &mut self.pending)
323                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
324            Self::saw(&mut self.first_byte, self.started, &mut self.parser, &self.pending);
325        }
326    }
327}
328
329impl Drop for UsageReader {
330    fn drop(&mut self) {
331        self.finish();
332    }
333}
rust_web_server/llm/usage.rs

rust_web_server/llm/
usage.rs
