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        if let Some(error) = chunk.error {
93            return Err(TransformError::with_source(
94                TransformErrorKind::InvalidResult,
95                "stream.error",
96                format!(
97                    "{} {}: {}",
98                    error.code,
99                    error.status.as_deref().unwrap_or_default(),
100                    error.message
101                ),
102                crate::transform::UpstreamError(error),
103            ));
104        }
105        merge_identity(&mut self.body.response_id, chunk.response_id, "response_id")?;
106        merge_identity(
107            &mut self.body.model_version,
108            chunk.model_version,
109            "model_version",
110        )?;
111        if let Some(usage) = chunk.usage_metadata {
112            let previous = self
113                .body
114                .usage_metadata
115                .get_or_insert_with(|| g::UsageMetadata::builder().build());
116            super::merge::usage(previous, usage)?;
117        }
118        if chunk.prompt_feedback.is_some() {
119            self.body.prompt_feedback = chunk.prompt_feedback;
120        }
121        if chunk.model_status.is_some() {
122            self.body.model_status = chunk.model_status;
123        }
124        if let Some(candidates) = chunk.candidates {
125            self.seen_candidates = true;
126            let multiple = candidates.len() > 1;
127            let mut in_chunk = BTreeSet::new();
128            for mut candidate in candidates {
129                let index = match candidate.index {
130                    Some(n) if n >= 0 => n,
131                    Some(_) => {
132                        return Err(TransformError::invalid_result(
133                            "candidate.index",
134                            "negative index",
135                        ));
136                    }
137                    None if !multiple => 0,
138                    None => {
139                        return Err(TransformError::missing_metadata(
140                            "multi-candidate stream index",
141                        ));
142                    }
143                };
144                if !in_chunk.insert(index) {
145                    return Err(TransformError::invalid_result(
146                        "candidate.index",
147                        "duplicate index within event",
148                    ));
149                }
150
151                candidate.index = Some(index);
152                let previous = self
153                    .candidates
154                    .entry(index)
155                    .or_insert_with(|| g::Candidate::builder().index(index).build());
156                if terminal(previous.finish_reason)
157                    && candidate
158                        .content
159                        .as_ref()
160                        .and_then(|c| c.parts.as_ref())
161                        .is_some_and(|p| !p.is_empty())
162                {
163                    return Err(TransformError::invalid_result(
164                        "candidate.content",
165                        "content after candidate terminal",
166                    ));
167                }
168                if let Some(content) = candidate.content.take() {
169                    let old = previous
170                        .content
171                        .get_or_insert_with(|| g::Content::builder().build());
172                    merge_identity(&mut old.role, content.role, "candidate.role")?;
173                    if let Some(parts) = content.parts {
174                        old.parts.get_or_insert_with(Vec::new).extend(parts);
175                    }
176                }
177                super::merge::candidate(previous, candidate)?;
178            }
179        }
180        Ok(())
181    }
182    /// EOF is successful only once every observed candidate has a terminal
183    /// reason, or the prompt was explicitly blocked before creating candidates.
184    pub fn finish(mut self) -> Result<Converted<g::GenerateContentResponseBody>, TransformError> {
185        if self.failed {
186            return Err(TransformError::invalid_result(
187                "stream",
188                "collector failed before EOF",
189            ));
190        }
191        let blocked = self
192            .body
193            .prompt_feedback
194            .as_ref()
195            .and_then(|v| v.block_reason)
196            .is_some_and(|r| r != g::BlockReason::Unspecified);
197        if self.bytes == 0
198            || (self.candidates.is_empty() && !blocked)
199            || self.candidates.values().any(|v| !terminal(v.finish_reason))
200        {
201            return Err(TransformError::invalid_result(
202                "stream",
203                "premature EOF without all candidate terminals",
204            ));
205        }
206        if self.seen_candidates {
207            self.body.candidates = Some(self.candidates.into_values().collect());
208        }
209        Ok(Converted {
210            value: self.body,
211            report: Default::default(),
212        })
213    }
214}
215
216pub(super) fn terminal(reason: Option<g::FinishReason>) -> bool {
217    reason.is_some_and(|r| r != g::FinishReason::Unspecified)
218}
219
220pub(super) fn limit(field: &str) -> TransformError {
221    TransformError::new(
222        TransformErrorKind::Limit,
223        format!("stream.{field}"),
224        "Gemini stream limit exceeded",
225    )
226}
227
228fn add_limit(total: &mut usize, n: usize, max: usize, field: &str) -> Result<(), TransformError> {
229    *total = total
230        .checked_add(n)
231        .filter(|n| *n <= max)
232        .ok_or_else(|| limit(field))?;
233    Ok(())
234}
235
236fn merge_identity(
237    old: &mut Option<String>,
238    new: Option<String>,
239    field: &str,
240) -> Result<(), TransformError> {
241    if let Some(value) = new {
242        if value.is_empty() || old.as_ref().is_some_and(|v| v != &value) {
243            return Err(TransformError::invalid_result(
244                field,
245                "empty or conflicting stream identity",
246            ));
247        }
248        *old = Some(value);
249    }
250    Ok(())
251}
gproxy_protocol/transform/generate/stream/gemini/collector.rs

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