gproxy_protocol/transform/generate/gemini_responses/stream/
common.rs

1use crate::{
2    transform::{Report, TransformError, TransformErrorKind, identity::IdentityFlow},
3    wire::DeclaredFields,
4};
5
6#[derive(Debug, Clone, Copy)]
7pub struct StreamLimits {
8    pub max_bytes: usize,
9    pub max_pending: usize,
10}
11
12impl Default for StreamLimits {
13    fn default() -> Self {
14        Self {
15            max_bytes: 16 * 1024 * 1024,
16            max_pending: 16 * 1024 * 1024,
17        }
18    }
19}
20
21#[derive(Debug)]
22pub struct StreamEnd<T> {
23    pub chunks: Vec<T>,
24    pub identities: IdentityFlow,
25    pub report: Report,
26}
27
28pub(super) fn invalid(message: impl Into<String>) -> TransformError {
29    TransformError::invalid_result("gemini_responses.stream", message)
30}
31
32pub(super) fn limit() -> TransformError {
33    TransformError::new(
34        TransformErrorKind::Limit,
35        "gemini_responses.stream",
36        "stream limit exceeded",
37    )
38}
39
40pub(super) struct Budget {
41    limits: StreamLimits,
42
43    input_bytes: usize,
44
45    output_bytes: usize,
46}
47
48impl Budget {
49    pub fn new(limits: StreamLimits) -> Self {
50        Self {
51            limits,
52
53            input_bytes: 0,
54
55            output_bytes: 0,
56        }
57    }
58    pub fn input<T: serde::Serialize>(&mut self, value: &T) -> Result<(), TransformError> {
59        self.input_bytes += measure(
60            value,
61            self.limits.max_bytes.saturating_sub(self.input_bytes),
62        )?;
63        Ok(())
64    }
65    pub fn output<T: serde::Serialize>(&mut self, value: &T) -> Result<usize, TransformError> {
66        let n = measure(
67            value,
68            self.limits.max_bytes.saturating_sub(self.output_bytes),
69        )?;
70        self.output_bytes += n;
71        Ok(n)
72    }
73}
74
75pub(super) fn measure<T: serde::Serialize>(value: &T, cap: usize) -> Result<usize, TransformError> {
76    let n = cap as u64;
77    crate::codec::encode_json(
78        value,
79        crate::codec::CodecLimits {
80            max_value_bytes: n,
81            max_body_bytes: n,
82            max_buffer_bytes: n,
83            max_line_bytes: n,
84            max_part_bytes: n,
85            max_parts: usize::MAX,
86        },
87    )
88    .map(|v| v.len())
89    .map_err(|e| {
90        if e.kind() == crate::codec::CodecErrorKind::Limit {
91            limit()
92        } else {
93            invalid(e.to_string())
94        }
95    })
96}
97
98pub(super) fn declared<T: DeclaredFields>(value: T) -> T {
99    value.into_declared()
100}
gproxy_protocol/transform/generate/gemini_responses/stream/common.rs

gproxy_protocol/transform/generate/gemini_responses/stream/
common.rs
