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

gproxy_protocol/usage/
mod.rs
