gproxy_protocol/usage/
mod.rs

1//! Per-call metering read from the standard response of an operation.
2//!
3//! A channel's job on the response side is to shape whatever its upstream
4//! answered into the standard form of the operation it served: Claude
5//! Messages, OpenAI Chat Completions or Responses, Gemini, an embedding or
6//! image reply. From there on every response is a standard one, whether it is
7//! converted for the client or passed through, so usage has one source: the
8//! readers here, chosen by operation and dialect. No reader knows a vendor.
9//! Fields a vendor adds beside the standard ones are the vendor's to read.
10//!
11//! Each reader has two forms. [`whole`] reads a complete response body.
12//! [`UsageReader`] watches a stream as it passes, chunk by chunk or message
13//! by message, without holding more of it than the event it is decoding:
14//! the stream is framed with this crate's own [`SseDecoder`] and
15//! [`JsonArrayDecoder`], events are filtered by name or by a substring before
16//! anything is parsed, and only the events that carry usage are deserialized
17//! — into small envelopes that skip the content around the usage object.
18//!
19//! A reading is `None` when the response reported no usage, which is not the
20//! same as zero consumption; a malformed response is one that reported none.
21//! No reader panics on its input. A streamed reading is `Partial` until the
22//! event that settles it has been seen, and whenever the stream was cut.
23//!
24//! What each operation reads:
25//!
26//! | operation | dialects | source |
27//! |---|---|---|
28//! | generate, stream generate, compact | Claude, OpenAI Chat, Responses (HTTP and websocket), Gemini | the dialect's usage object |
29//! | count tokens | all | only a usage object the upstream attaches; the count itself was never consumed |
30//! | embeddings | OpenAI, Gemini | `usage` / `usageMetadata`, input only |
31//! | images | OpenAI | `usage`, the number of images, and the image qualifiers |
32//! | transcription, translation | OpenAI | `usage` in tokens or seconds |
33//! | rerank | OpenAI | `usage` in tokens or search units |
34//! | realtime | OpenAI | each `response.done`, summed by response id |
35//!
36//! [`SseDecoder`]: crate::codec::SseDecoder
37//! [`JsonArrayDecoder`]: crate::codec::JsonArrayDecoder
38
39mod claude;
40mod common;
41mod gemini;
42mod media;
43mod openai;
44mod types;
45
46pub use types::{
47    NormalizedUsage, ResponseUsage, TokenUsage, UsageAttempt, UsageCompleteness, UsageStreamEnd,
48    UsageTransport,
49};
50
51use crate::codec::{CodecLimits, JsonArrayDecoder, SseDecoder, SseFrame};
52use crate::connection::StreamFraming;
53use crate::{Dialect, Operation, WireFamily};
54
55/// What a response of an operation in a dialect is read as.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57enum Reader {
58    Claude,
59    Chat,
60    Responses,
61    Gemini,
62    Realtime,
63    OpenAiEmbedding,
64    GeminiEmbedding,
65    Image,
66    Transcription,
67    Rerank,
68    /// Count tokens: only a usage object the upstream attached, read the way
69    /// the dialect's generation is.
70    CountTokens(&'static Reader),
71}
72
73fn reader(operation: Operation, dialect: Dialect) -> Option<Reader> {
74    let generate = match dialect {
75        Dialect::Claude => Reader::Claude,
76        Dialect::OpenAiChat => Reader::Chat,
77        Dialect::OpenAi | Dialect::OpenAiResponsesWebSocket => Reader::Responses,
78        Dialect::Gemini => Reader::Gemini,
79    };
80    let openai = dialect.family() == WireFamily::OpenAi;
81    Some(match operation {
82        Operation::GenerateContent
83        | Operation::StreamGenerateContent
84        | Operation::CompactContent => generate,
85        Operation::CountTokens => Reader::CountTokens(match generate {
86            Reader::Claude => &Reader::Claude,
87            Reader::Chat => &Reader::Chat,
88            Reader::Gemini => &Reader::Gemini,
89            _ => &Reader::Responses,
90        }),
91        Operation::CreateEmbedding | Operation::BatchCreateEmbedding => match dialect.family() {
92            WireFamily::OpenAi => Reader::OpenAiEmbedding,
93            WireFamily::Gemini => Reader::GeminiEmbedding,
94            WireFamily::Claude => return None,
95        },
96        Operation::CreateImage | Operation::EditImage if openai => Reader::Image,
97        Operation::CreateTranscription | Operation::CreateTranslation if openai => {
98            Reader::Transcription
99        }
100        Operation::Rerank if openai => Reader::Rerank,
101        Operation::ConnectRealtime if openai => Reader::Realtime,
102        _ => return None,
103    })
104}
105
106/// Usage from a complete, successful response body of `operation` in
107/// `dialect`. The caller decides whether the response succeeded; an error
108/// body reports no usage anyway.
109pub fn whole(operation: Operation, dialect: Dialect, body: &[u8]) -> Option<NormalizedUsage> {
110    read(reader(operation, dialect)?, body)
111}
112
113/// Whether a complete response of `operation` in `dialect` has usage for
114/// [`whole`] to read. A caller that must hold a body to read it whole holds
115/// it only when this says there is something to find.
116pub fn reads(operation: Operation, dialect: Dialect) -> bool {
117    reader(operation, dialect).is_some()
118}
119
120fn read(reader: Reader, body: &[u8]) -> Option<NormalizedUsage> {
121    match reader {
122        Reader::Claude => claude::whole(body),
123        Reader::Chat | Reader::Responses | Reader::Realtime => openai::whole(body),
124        Reader::Gemini => gemini::whole(body),
125        Reader::OpenAiEmbedding => media::openai_embedding(body),
126        Reader::GeminiEmbedding => gemini::embedding(body),
127        Reader::Image => media::image(body),
128        Reader::Transcription => media::transcription(body),
129        Reader::Rerank => media::rerank(body),
130        Reader::CountTokens(generate) => read(*generate, body),
131    }
132}
133
134/// Bounds for watching a stream. The host enforces the real transfer limits;
135/// these only keep the reader's own buffer finite. They are generous because
136/// the event that carries usage can also carry the whole answer: a Responses
137/// `response.completed` repeats every output item, base64 images included,
138/// and an event over the bound would lose the usage with it.
139const LIMITS: CodecLimits = CodecLimits {
140    max_buffer_bytes: 64 * 1024 * 1024,
141    max_value_bytes: 64 * 1024 * 1024,
142    max_body_bytes: u64::MAX,
143    max_line_bytes: 64 * 1024 * 1024,
144    max_part_bytes: 0,
145    max_parts: 0,
146};
147
148/// How the stream arrives.
149enum Framing {
150    Sse(SseDecoder),
151    JsonArray(JsonArrayDecoder),
152    /// Websocket messages, one complete event each.
153    Messages,
154    /// The framing broke; the stream is no longer read, and what was read
155    /// before stays as a partial reading.
156    Broken,
157}
158
159/// The per-dialect state of a watched stream.
160enum Watch {
161    Claude(claude::Stream),
162    Chat(openai::ChatStream),
163    Responses(openai::ResponsesStream),
164    Gemini(gemini::Stream),
165    Transcription(media::TranscriptionStream),
166}
167
168impl Watch {
169    /// Whether the event carried usage that the reading now reflects.
170    fn event(&mut self, name: Option<&str>, data: &str) -> bool {
171        match self {
172            Self::Claude(stream) => stream.event(name, data),
173            Self::Chat(stream) => stream.event(data),
174            Self::Responses(stream) => stream.event(name, data),
175            Self::Gemini(stream) => stream.event(data),
176            Self::Transcription(stream) => stream.event(name, data),
177        }
178    }
179
180    /// The reading so far. `settled` says the stream ended on its own,
181    /// which only Gemini's reading depends on.
182    fn reading(&self, settled: bool) -> Option<NormalizedUsage> {
183        match self {
184            Self::Claude(stream) => stream.snapshot(),
185            Self::Chat(stream) => stream.snapshot(),
186            Self::Responses(stream) => stream.snapshot(),
187            Self::Gemini(stream) => stream.reading(settled),
188            Self::Transcription(stream) => stream.snapshot(),
189        }
190    }
191}
192
193/// Watches one streamed response for its usage.
194///
195/// Feed it the stream as it passes — raw HTTP chunks with [`push`], or
196/// websocket text messages with [`push_message`] — and take the reading with
197/// [`finish`]. It borrows what it is given and never alters the stream.
198/// Every reading is cumulative: [`snapshot`] and [`finish`] replace earlier
199/// readings rather than adding to them.
200///
201/// [`push`]: UsageReader::push
202/// [`push_message`]: UsageReader::push_message
203/// [`snapshot`]: UsageReader::snapshot
204/// [`finish`]: UsageReader::finish
205pub struct UsageReader {
206    framing: Framing,
207    watch: Watch,
208    /// Whether to keep the latest event that carried usage; see
209    /// [`UsageReader::keep_usage_event`].
210    keep: bool,
211    kept: Option<Kept>,
212}
213
214/// The latest event that carried usage, as it arrived.
215enum Kept {
216    Text(String),
217    Value(serde_json::Value),
218}
219
220impl UsageReader {
221    /// A reader for a stream of `operation` in `dialect` arriving over
222    /// `transport`, or `None` when that stream reports no usage to watch for:
223    /// an operation without a streamed form, or a transport its dialect does
224    /// not stream over. An HTTP stream without a declared framing is read the
225    /// way the dialect streams by default: a JSON array for Gemini, SSE for
226    /// the rest.
227    pub fn new(operation: Operation, dialect: Dialect, transport: UsageTransport) -> Option<Self> {
228        let reader = reader(operation, dialect)?;
229        let watch = match reader {
230            Reader::Claude => Watch::Claude(claude::Stream::default()),
231            Reader::Chat => Watch::Chat(openai::ChatStream::default()),
232            Reader::Responses => {
233                Watch::Responses(openai::ResponsesStream::new(openai::Mode::Single))
234            }
235            Reader::Realtime => {
236                Watch::Responses(openai::ResponsesStream::new(openai::Mode::Realtime))
237            }
238            Reader::Image => Watch::Responses(openai::ResponsesStream::new(openai::Mode::Image)),
239            Reader::Gemini => Watch::Gemini(gemini::Stream::default()),
240            Reader::Transcription => Watch::Transcription(media::TranscriptionStream::default()),
241            Reader::OpenAiEmbedding
242            | Reader::GeminiEmbedding
243            | Reader::Rerank
244            | Reader::CountTokens(_) => return None,
245        };
246        let framing = match (transport, &watch) {
247            (UsageTransport::WebSocket, Watch::Responses(_)) => Framing::Messages,
248            (UsageTransport::WebSocket, _) => return None,
249            (
250                UsageTransport::Http {
251                    framing: Some(StreamFraming::JsonArray) | None,
252                },
253                Watch::Gemini(_),
254            ) => Framing::JsonArray(JsonArrayDecoder::new(LIMITS)),
255            (
256                UsageTransport::Http {
257                    framing: Some(StreamFraming::Sse) | None,
258                },
259                _,
260            ) => Framing::Sse(SseDecoder::new(LIMITS)),
261            (UsageTransport::Http { .. }, _) => return None,
262        };
263        Some(Self {
264            framing,
265            watch,
266            keep: false,
267            kept: None,
268        })
269    }
270
271    /// Also keep the latest event that carried usage, for
272    /// [`usage_event`](UsageReader::usage_event). A vendor may report fields
273    /// of its own beside the standard usage object, such as a price it
274    /// charged; those are not this module's to read, and whoever reads them
275    /// needs the event they arrived in. Off by default, because the kept
276    /// event is a copy: a Responses `response.completed` repeats the whole
277    /// answer.
278    pub fn keep_usage_event(&mut self) {
279        self.keep = true;
280    }
281
282    /// The latest event that carried usage, when
283    /// [`keep_usage_event`](UsageReader::keep_usage_event) asked for it: the
284    /// Claude `message_delta` (or the `message_start` before any delta), the
285    /// Chat chunk with a usage object, the Responses terminal event, the
286    /// Gemini chunk with `usageMetadata`. `None` when no event carried usage.
287    pub fn usage_event(&self) -> Option<serde_json::Value> {
288        match self.kept.as_ref()? {
289            Kept::Text(text) => serde_json::from_str(text).ok(),
290            Kept::Value(value) => Some(value.clone()),
291        }
292    }
293
294    /// Observe raw HTTP bytes. Chunks need not align with events.
295    pub fn push(&mut self, chunk: &[u8]) {
296        match &mut self.framing {
297            Framing::Sse(decoder) => match decoder.push(chunk) {
298                Ok(frames) => {
299                    for frame in frames {
300                        if let SseFrame::Event(event) = frame
301                            && self.watch.event(event.event.as_deref(), &event.data)
302                            && self.keep
303                        {
304                            self.kept = Some(Kept::Text(event.data));
305                        }
306                    }
307                }
308                Err(_) => self.framing = Framing::Broken,
309            },
310            Framing::JsonArray(decoder) => match decoder.push(chunk) {
311                Ok(records) => {
312                    if let Watch::Gemini(stream) = &mut self.watch {
313                        for record in records {
314                            if stream.record(&record) && self.keep {
315                                self.kept = Some(Kept::Value(record));
316                            }
317                        }
318                    }
319                }
320                Err(_) => self.framing = Framing::Broken,
321            },
322            Framing::Messages | Framing::Broken => {}
323        }
324    }
325
326    /// Observe one websocket text message, which is one complete event.
327    pub fn push_message(&mut self, text: &str) {
328        if matches!(self.framing, Framing::Messages) && self.watch.event(None, text) && self.keep {
329            self.kept = Some(Kept::Text(text.to_owned()));
330        }
331    }
332
333    /// The reading so far. It is `Partial` until the event that settles it
334    /// has been seen.
335    pub fn snapshot(&self) -> Option<NormalizedUsage> {
336        self.reading(false)
337    }
338
339    fn reading(&self, settled: bool) -> Option<NormalizedUsage> {
340        let broken = matches!(self.framing, Framing::Broken);
341        let mut usage = self.watch.reading(settled && !broken)?;
342        if broken {
343            usage.completeness = UsageCompleteness::Partial;
344        }
345        Some(usage)
346    }
347
348    /// The final reading. A stream that was cut is never complete, whatever
349    /// it reported before it stopped.
350    pub fn finish(self, end: UsageStreamEnd) -> Option<NormalizedUsage> {
351        let mut usage = self.reading(end == UsageStreamEnd::Complete)?;
352        if end == UsageStreamEnd::Interrupted {
353            usage.completeness = UsageCompleteness::Partial;
354        }
355        Some(usage)
356    }
357}
358
359#[cfg(test)]
360pub(crate) mod tests {
361    use super::*;
362
363    /// `wire` in chunks of `size` bytes, splitting inside characters too.
364    pub(crate) fn chunked(wire: &str, size: usize) -> Vec<Vec<u8>> {
365        wire.as_bytes()
366            .chunks(size.clamp(1, wire.len().max(1)))
367            .map(<[u8]>::to_vec)
368            .collect()
369    }
370
371    pub(crate) fn feed_chunks(
372        operation: Operation,
373        dialect: Dialect,
374        transport: UsageTransport,
375        chunks: &[Vec<u8>],
376        end: UsageStreamEnd,
377    ) -> Option<NormalizedUsage> {
378        let mut reader = UsageReader::new(operation, dialect, transport).expect("watchable");
379        for chunk in chunks {
380            reader.push(chunk);
381        }
382        reader.finish(end)
383    }
384
385    pub(crate) fn feed_sse(
386        operation: Operation,
387        dialect: Dialect,
388        chunks: &[Vec<u8>],
389        end: UsageStreamEnd,
390    ) -> Option<NormalizedUsage> {
391        feed_chunks(
392            operation,
393            dialect,
394            UsageTransport::Http {
395                framing: Some(StreamFraming::Sse),
396            },
397            chunks,
398            end,
399        )
400    }
401
402    pub(crate) fn feed_messages(
403        operation: Operation,
404        dialect: Dialect,
405        messages: &[&str],
406        end: UsageStreamEnd,
407    ) -> Option<NormalizedUsage> {
408        let mut reader =
409            UsageReader::new(operation, dialect, UsageTransport::WebSocket).expect("watchable");
410        for message in messages {
411            reader.push_message(message);
412        }
413        reader.finish(end)
414    }
415
416    #[test]
417    fn count_tokens_never_meters_the_count_itself() {
418        for (dialect, body) in [
419            (Dialect::Claude, r#"{"input_tokens":12}"#),
420            (
421                Dialect::OpenAi,
422                r#"{"object":"response.input_tokens","input_tokens":12}"#,
423            ),
424            (
425                Dialect::Gemini,
426                r#"{"totalTokens":12,"cachedContentTokenCount":2}"#,
427            ),
428        ] {
429            assert!(
430                whole(Operation::CountTokens, dialect, body.as_bytes()).is_none(),
431                "{dialect:?}"
432            );
433        }
434        let attached = whole(
435            Operation::CountTokens,
436            Dialect::Claude,
437            br#"{"input_tokens":12,"usage":{"input_tokens":3,"output_tokens":0}}"#,
438        )
439        .expect("an attached usage object is reported");
440        assert_eq!(attached.tokens.input_tokens, Some(3));
441    }
442
443    #[test]
444    fn only_operations_with_a_stream_get_a_stream_reader() {
445        let sse = UsageTransport::Http {
446            framing: Some(StreamFraming::Sse),
447        };
448        for (operation, dialect) in [
449            (Operation::CountTokens, Dialect::Claude),
450            (Operation::CreateEmbedding, Dialect::OpenAi),
451            (Operation::Rerank, Dialect::OpenAi),
452            (Operation::ListModels, Dialect::OpenAi),
453            (Operation::DeleteFile, Dialect::OpenAi),
454            (Operation::CreateImage, Dialect::Gemini),
455        ] {
456            assert!(
457                UsageReader::new(operation, dialect, sse).is_none(),
458                "{operation:?}"
459            );
460        }
461        assert!(
462            UsageReader::new(
463                Operation::StreamGenerateContent,
464                Dialect::Claude,
465                UsageTransport::WebSocket
466            )
467            .is_none()
468        );
469        assert!(
470            UsageReader::new(
471                Operation::StreamGenerateContent,
472                Dialect::OpenAiChat,
473                UsageTransport::Http {
474                    framing: Some(StreamFraming::NdJson)
475                }
476            )
477            .is_none()
478        );
479    }
480
481    /// The kept event is the last one that carried usage, whatever came
482    /// after it, and nothing is kept unless asked for.
483    #[test]
484    fn the_latest_usage_event_is_kept_on_request() {
485        let sse = UsageTransport::Http {
486            framing: Some(StreamFraming::Sse),
487        };
488        let claude = concat!(
489            "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"m\",\"usage\":{\"input_tokens\":3,\"output_tokens\":1}}}\n\n",
490            "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{},\"usage\":{\"output_tokens\":9},\"cost\":2}\n\n",
491            "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
492        );
493        let mut reader =
494            UsageReader::new(Operation::StreamGenerateContent, Dialect::Claude, sse).unwrap();
495        reader.push(claude.as_bytes());
496        assert!(reader.usage_event().is_none());
497        let mut reader =
498            UsageReader::new(Operation::StreamGenerateContent, Dialect::Claude, sse).unwrap();
499        reader.keep_usage_event();
500        for chunk in chunked(claude, 5) {
501            reader.push(&chunk);
502        }
503        let event = reader.usage_event().unwrap();
504        assert_eq!(event["type"], "message_delta");
505        assert_eq!(event["cost"], 2);
506
507        let chat = "data: {\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":1},\"model\":\"x\"}\n\ndata: {\"usage\":null}\n\ndata: [DONE]\n\n";
508        let mut reader =
509            UsageReader::new(Operation::StreamGenerateContent, Dialect::OpenAiChat, sse).unwrap();
510        reader.keep_usage_event();
511        reader.push(chat.as_bytes());
512        assert_eq!(reader.usage_event().unwrap()["model"], "x");
513
514        let gemini = r#"[{"candidates":[]},{"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2},"modelVersion":"g"},{"candidates":[]}]"#;
515        let mut reader = UsageReader::new(
516            Operation::StreamGenerateContent,
517            Dialect::Gemini,
518            UsageTransport::Http { framing: None },
519        )
520        .unwrap();
521        reader.keep_usage_event();
522        reader.push(gemini.as_bytes());
523        assert_eq!(reader.usage_event().unwrap()["modelVersion"], "g");
524
525        let mut reader = UsageReader::new(
526            Operation::StreamGenerateContent,
527            Dialect::OpenAiResponsesWebSocket,
528            UsageTransport::WebSocket,
529        )
530        .unwrap();
531        reader.keep_usage_event();
532        reader.push_message(r#"{"type":"response.created","response":{"id":"r"}}"#);
533        assert!(reader.usage_event().is_none());
534        reader.push_message(r#"{"type":"response.completed","response":{"id":"r","usage":{"input_tokens":1,"output_tokens":1}}}"#);
535        assert_eq!(reader.usage_event().unwrap()["response"]["id"], "r");
536    }
537
538    #[test]
539    fn a_broken_stream_keeps_its_reading_as_partial() {
540        let usage = "data: {\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":1}}\n\n";
541        let mut reader = UsageReader::new(
542            Operation::StreamGenerateContent,
543            Dialect::OpenAiChat,
544            UsageTransport::Http {
545                framing: Some(StreamFraming::Sse),
546            },
547        )
548        .unwrap();
549        reader.push(usage.as_bytes());
550        reader.push(&[0xff, 0xfe, b'\n']);
551        reader.push(usage.as_bytes());
552        let reading = reader.finish(UsageStreamEnd::Complete).unwrap();
553        assert_eq!(reading.tokens.input_tokens, Some(2));
554        assert_eq!(reading.completeness, UsageCompleteness::Partial);
555    }
556
557    #[test]
558    fn arbitrary_bytes_never_panic() {
559        let mut seed = 0x9e37_79b9_u32;
560        let mut noise = Vec::new();
561        for _ in 0..4096 {
562            seed ^= seed << 13;
563            seed ^= seed >> 17;
564            seed ^= seed << 5;
565            noise.push(b"{}[]\":,data: event\n\r0123usage_metadata"[seed as usize % 38]);
566        }
567        for dialect in [
568            Dialect::Claude,
569            Dialect::OpenAiChat,
570            Dialect::OpenAi,
571            Dialect::Gemini,
572            Dialect::OpenAiResponsesWebSocket,
573        ] {
574            for operation in [
575                Operation::GenerateContent,
576                Operation::StreamGenerateContent,
577                Operation::CountTokens,
578                Operation::CreateEmbedding,
579                Operation::CreateImage,
580                Operation::CreateTranscription,
581                Operation::Rerank,
582                Operation::ConnectRealtime,
583            ] {
584                let _ = whole(operation, dialect, &noise);
585                for transport in [
586                    UsageTransport::Http { framing: None },
587                    UsageTransport::Http {
588                        framing: Some(StreamFraming::JsonArray),
589                    },
590                    UsageTransport::WebSocket,
591                ] {
592                    if let Some(mut reader) = UsageReader::new(operation, dialect, transport) {
593                        for chunk in noise.chunks(13) {
594                            reader.push(chunk);
595                            reader.push_message(&String::from_utf8_lossy(chunk));
596                        }
597                        let _ = reader.finish(UsageStreamEnd::Interrupted);
598                    }
599                }
600            }
601        }
602    }
603}
gproxy_protocol/usage/mod.rs

gproxy_protocol/usage/
mod.rs
