gproxy_protocol/transform/generate/stream/gemini/
collector.rs

1use crate::{
2    codec::{self, CodecErrorKind, CodecLimits},
3    transform::{Converted, TransformError, TransformErrorKind},
4    wire::{DeclaredFields, gemini as g},
5};
6use std::collections::{BTreeMap, BTreeSet};
7
8#[derive(Debug, Clone, Copy)]
9pub struct GeminiStreamLimits {
10    pub max_bytes: usize,
11}
12
13impl Default for GeminiStreamLimits {
14    fn default() -> Self {
15        Self {
16            max_bytes: 16 * 1024 * 1024,
17        }
18    }
19}
20
21/// Collects native chunks without retaining a second copy of the event log.
22/// Content is incremental; usage and non-content metadata are native cumulative
23/// snapshots. A failed push poisons the collector so partial state cannot pass.
24pub struct GeminiStreamCollector {
25    limits: GeminiStreamLimits,
26
27    bytes: usize,
28
29    failed: bool,
30    body: g::GenerateContentResponseBody,
31    candidates: BTreeMap<i64, g::Candidate>,
32    seen_candidates: bool,
33}
34
35impl GeminiStreamCollector {
36    pub fn new(limits: GeminiStreamLimits) -> Self {
37        Self {
38            limits,
39
40            bytes: 0,
41
42            failed: false,
43            body: g::GenerateContentResponseBody::builder().build(),
44            candidates: BTreeMap::new(),
45            seen_candidates: false,
46        }
47    }
48    pub fn push(&mut self, chunk: g::GenerateContentResponseBody) -> Result<(), TransformError> {
49        if self.failed {
50            return Err(TransformError::invalid_result(
51                "stream",
52                "collector already failed",
53            ));
54        }
55        let result = self.push_declared(chunk.into_declared());
56        if result.is_err() {
57            self.failed = true;
58        }
59        result
60    }
61    fn push_declared(
62        &mut self,
63        chunk: g::GenerateContentResponseBody,
64    ) -> Result<(), TransformError> {
65        // Serialization measures the already rebuilt declared event. It is
66        // never used to discover fields or convert between protocols.
67        let remaining = self.limits.max_bytes.saturating_sub(self.bytes) as u64;
68        let encoded = codec::encode_json(
69            &chunk,
70            CodecLimits {
71                max_buffer_bytes: remaining,
72                max_value_bytes: remaining,
73                max_body_bytes: remaining,
74                max_line_bytes: remaining,
75                max_part_bytes: remaining,
76                max_parts: 0,
77            },
78        )
79        .map_err(|error| {
80            if error.kind() == CodecErrorKind::Limit {
81                limit("bytes")
82            } else {
83                TransformError::invalid_result("stream.event", error.to_string())
84            }
85        })?;
86        add_limit(
87            &mut self.bytes,
88            encoded.len(),
89            self.limits.max_bytes,
90            "bytes",
91        )?;
92        merge_identity(&mut self.body.response_id, chunk.response_id, "response_id")?;
93        merge_identity(
94            &mut self.body.model_version,
95            chunk.model_version,
96            "model_version",
97        )?;
98        if let Some(usage) = chunk.usage_metadata {
99            let previous = self
100                .body
101                .usage_metadata
102                .get_or_insert_with(|| g::UsageMetadata::builder().build());
103            super::merge::usage(previous, usage)?;
104        }
105        if chunk.prompt_feedback.is_some() {
106            self.body.prompt_feedback = chunk.prompt_feedback;
107        }
108        if chunk.model_status.is_some() {
109            self.body.model_status = chunk.model_status;
110        }
111        if let Some(candidates) = chunk.candidates {
112            self.seen_candidates = true;
113            let multiple = candidates.len() > 1;
114            let mut in_chunk = BTreeSet::new();
115            for mut candidate in candidates {
116                let index = match candidate.index {
117                    Some(n) if n >= 0 => n,
118                    Some(_) => {
119                        return Err(TransformError::invalid_result(
120                            "candidate.index",
121                            "negative index",
122                        ));
123                    }
124                    None if !multiple => 0,
125                    None => {
126                        return Err(TransformError::missing_metadata(
127                            "multi-candidate stream index",
128                        ));
129                    }
130                };
131                if !in_chunk.insert(index) {
132                    return Err(TransformError::invalid_result(
133                        "candidate.index",
134                        "duplicate index within event",
135                    ));
136                }
137
138                candidate.index = Some(index);
139                let previous = self
140                    .candidates
141                    .entry(index)
142                    .or_insert_with(|| g::Candidate::builder().index(index).build());
143                if terminal(previous.finish_reason)
144                    && candidate
145                        .content
146                        .as_ref()
147                        .and_then(|c| c.parts.as_ref())
148                        .is_some_and(|p| !p.is_empty())
149                {
150                    return Err(TransformError::invalid_result(
151                        "candidate.content",
152                        "content after candidate terminal",
153                    ));
154                }
155                if let Some(content) = candidate.content.take() {
156                    let old = previous
157                        .content
158                        .get_or_insert_with(|| g::Content::builder().build());
159                    merge_identity(&mut old.role, content.role, "candidate.role")?;
160                    if let Some(parts) = content.parts {
161                        old.parts.get_or_insert_with(Vec::new).extend(parts);
162                    }
163                }
164                super::merge::candidate(previous, candidate)?;
165            }
166        }
167        Ok(())
168    }
169    /// EOF is successful only once every observed candidate has a terminal
170    /// reason, or the prompt was explicitly blocked before creating candidates.
171    pub fn finish(mut self) -> Result<Converted<g::GenerateContentResponseBody>, TransformError> {
172        if self.failed {
173            return Err(TransformError::invalid_result(
174                "stream",
175                "collector failed before EOF",
176            ));
177        }
178        let blocked = self
179            .body
180            .prompt_feedback
181            .as_ref()
182            .and_then(|v| v.block_reason)
183            .is_some_and(|r| r != g::BlockReason::Unspecified);
184        if self.bytes == 0
185            || (self.candidates.is_empty() && !blocked)
186            || self.candidates.values().any(|v| !terminal(v.finish_reason))
187        {
188            return Err(TransformError::invalid_result(
189                "stream",
190                "premature EOF without all candidate terminals",
191            ));
192        }
193        if self.seen_candidates {
194            self.body.candidates = Some(self.candidates.into_values().collect());
195        }
196        Ok(Converted {
197            value: self.body,
198            report: Default::default(),
199        })
200    }
201}
202
203pub(super) fn terminal(reason: Option<g::FinishReason>) -> bool {
204    reason.is_some_and(|r| r != g::FinishReason::Unspecified)
205}
206
207pub(super) fn limit(field: &str) -> TransformError {
208    TransformError::new(
209        TransformErrorKind::Limit,
210        format!("stream.{field}"),
211        "Gemini stream limit exceeded",
212    )
213}
214
215fn add_limit(total: &mut usize, n: usize, max: usize, field: &str) -> Result<(), TransformError> {
216    *total = total
217        .checked_add(n)
218        .filter(|n| *n <= max)
219        .ok_or_else(|| limit(field))?;
220    Ok(())
221}
222
223fn merge_identity(
224    old: &mut Option<String>,
225    new: Option<String>,
226    field: &str,
227) -> Result<(), TransformError> {
228    if let Some(value) = new {
229        if value.is_empty() || old.as_ref().is_some_and(|v| v != &value) {
230            return Err(TransformError::invalid_result(
231                field,
232                "empty or conflicting stream identity",
233            ));
234        }
235        *old = Some(value);
236    }
237    Ok(())
238}
gproxy_protocol/transform/generate/stream/gemini/collector.rs

gproxy_protocol/transform/generate/stream/gemini/
collector.rs
