gproxy_protocol/transform/generate/stream/responses/
lifecycle.rs

1use super::{
2    items::{Item, item_id},
3    parts::Part,
4    *,
5};
6use crate::transform::Converted;
7
8impl ResponsesStreamCollector {
9    pub(super) fn check_response(
10        &self,
11        v: &r::GenerateContentResponseBody,
12        terminal: bool,
13    ) -> Result<(), TransformError> {
14        if v.id.is_empty() || v.model.is_empty() || v.created_at < 0 {
15            return Err(invalid("missing/invalid native response identity facts"));
16        }
17        if let Some(old) = &self.response
18            && (old.id != v.id || old.model != v.model || old.created_at != v.created_at)
19        {
20            return Err(invalid("response identity/model/creation changed"));
21        }
22        if !terminal
23            && (!matches!(
24                v.status,
25                Some(r::ResponseStatus::InProgress | r::ResponseStatus::Queued)
26            ) || v.error.is_some()
27                || v.incomplete_details.is_some()
28                || v.completed_at.flatten().is_some())
29        {
30            return Err(invalid(
31                "nonterminal response contains terminal status/facts",
32            ));
33        }
34        Ok(())
35    }
36    pub(super) fn terminal(
37        &mut self,
38        v: r::GenerateContentResponseBody,
39        status: r::ResponseStatus,
40    ) -> Result<(), TransformError> {
41        self.check_response(&v, true)?;
42        for item in &v.output {
43            super::status::item(item, status == r::ResponseStatus::Completed)?;
44        }
45        let (text, json) = v
46            .output
47            .iter()
48            .map(item_sizes)
49            .fold((0usize, 0usize), |(a, b), (c, d)| {
50                (a.saturating_add(c), b.saturating_add(d))
51            });
52        if text > self.limits.max_text_bytes || json > self.limits.max_json_bytes {
53            return Err(limit());
54        }
55
56        if v.status != Some(status)
57            || v.error.is_some()
58            || (status == r::ResponseStatus::Completed && v.incomplete_details.is_some())
59        {
60            return Err(invalid("terminal response status mismatch"));
61        }
62        if status == r::ResponseStatus::Incomplete
63            && v.incomplete_details
64                .as_ref()
65                .and_then(|v| v.reason)
66                .is_none()
67        {
68            return Err(invalid("incomplete response missing reason"));
69        }
70        if status == r::ResponseStatus::Completed
71            && self
72                .items
73                .iter()
74                .any(|item| item.parts.values().any(|part| part.incomplete))
75        {
76            return Err(invalid(
77                "completed response contains incomplete summary part",
78            ));
79        }
80        if self.items.len() != v.output.len()
81            || self
82                .items
83                .iter()
84                .zip(&v.output)
85                .any(|(a, b)| !a.done || &a.value != b)
86        {
87            return Err(invalid(
88                "terminal response contradicts/misses completed output items",
89            ));
90        }
91        if let Some(Some(text)) = &v.output_text {
92            let actual: String = v
93                .output
94                .iter()
95                .filter_map(|v| {
96                    if let r::ResponseOutputItem::Message(v) = v {
97                        Some(v)
98                    } else {
99                        None
100                    }
101                })
102                .flat_map(|v| &v.content)
103                .filter_map(|v| {
104                    if let i::OutputContent::Text(v) = v {
105                        Some(v.text.as_str())
106                    } else {
107                        None
108                    }
109                })
110                .collect();
111            if text != &actual {
112                return Err(invalid("output_text contradicts output"));
113            }
114        }
115        if let Some(u) = v.usage.as_ref().and_then(Option::as_ref)
116            && ([
117                u.input_tokens,
118                u.output_tokens,
119                u.total_tokens,
120                u.input_tokens_details.cached_tokens,
121                u.input_tokens_details.cache_write_tokens,
122                u.output_tokens_details.reasoning_tokens,
123            ]
124            .iter()
125            .any(|v| *v < 0)
126                || u.input_tokens.checked_add(u.output_tokens) != Some(u.total_tokens)
127                || u.input_tokens_details
128                    .cached_tokens
129                    .checked_add(u.input_tokens_details.cache_write_tokens)
130                    .is_none_or(|total| total > u.input_tokens)
131                || u.output_tokens_details.reasoning_tokens > u.output_tokens)
132        {
133            return Err(invalid("invalid terminal usage"));
134        }
135        self.response = Some(v);
136        self.done = true;
137        Ok(())
138    }
139    pub(super) fn item(&mut self, index: i64, id: &str) -> Result<&mut Item, TransformError> {
140        let item = self
141            .items
142            .get_mut(usize::try_from(index).map_err(|_| invalid("negative output index"))?)
143            .ok_or_else(|| invalid("unknown output index"))?;
144        if !item.done && item_id(&item.value).is_none() {
145            if id.is_empty() || !self.ids.insert(id.to_owned()) {
146                return Err(invalid("empty/duplicate late item ID"));
147            }
148            match &mut item.value {
149                r::ResponseOutputItem::FunctionCall(v) => v.id = Some(id.into()),
150                r::ResponseOutputItem::CustomToolCall(v) => v.id = Some(id.into()),
151                _ => return Err(invalid("item cannot receive late ID")),
152            }
153        }
154        if item.done || item_id(&item.value) != Some(id) {
155            return Err(invalid("item ID/index/phase mismatch"));
156        }
157        Ok(item)
158    }
159    pub(super) fn part(
160        &mut self,
161        index: i64,
162        id: &str,
163        content: i64,
164        summary: bool,
165    ) -> Result<&mut Part, TransformError> {
166        self.item(index, id)?
167            .parts
168            .get_mut(&(summary, content))
169            .ok_or_else(|| invalid("content part not added"))
170    }
171    pub(super) fn add_part(
172        &mut self,
173        index: i64,
174        id: &str,
175        content: i64,
176        summary: bool,
177        value: s::OutputContentPart,
178    ) -> Result<(), TransformError> {
179        let text = match &value {
180            s::OutputContentPart::Text(v) => &v.text,
181            s::OutputContentPart::Refusal(v) => &v.refusal,
182            s::OutputContentPart::Reasoning(v) => &v.text,
183        };
184        self.charge(text.len(), false)?;
185        let item = self.item(index, id)?;
186        let expected =
187            item.base_parts(summary)? + item.parts.keys().filter(|(s, _)| *s == summary).count();
188        if content != expected as i64 {
189            return Err(invalid("noncontiguous/duplicate content index"));
190        }
191        if !matches!(
192            (&item.value, &value, summary),
193            (
194                r::ResponseOutputItem::Message(_),
195                s::OutputContentPart::Text(_) | s::OutputContentPart::Refusal(_),
196                false
197            ) | (
198                r::ResponseOutputItem::Reasoning(_),
199                s::OutputContentPart::Reasoning(_),
200                _
201            )
202        ) {
203            return Err(invalid("content type incompatible with item"));
204        }
205        item.parts
206            .insert((summary, content), Part::new(value, summary));
207        Ok(())
208    }
209    pub(super) fn delta(
210        &mut self,
211        index: i64,
212        id: &str,
213        content: i64,
214        kind: u8,
215        text: &str,
216    ) -> Result<&mut Part, TransformError> {
217        self.charge(text.len(), false)?;
218        let p = self.part(index, id, content, kind == 3)?;
219        p.check_kind(kind)?;
220        p.seeded = false;
221        p.text_mut().push_str(text);
222        Ok(p)
223    }
224    pub(super) fn text_done(
225        &mut self,
226        index: i64,
227        id: &str,
228        content: i64,
229        kind: u8,
230        text: &str,
231        logs: &[s::StreamLogprob],
232    ) -> Result<(), TransformError> {
233        let p = self.part(index, id, content, kind == 3)?;
234        p.check_kind(kind)?;
235        p.finish_text(text, logs)
236    }
237    pub(super) fn argument(
238        &mut self,
239        index: i64,
240        id: &str,
241        kind: u8,
242        text: &str,
243        done: bool,
244    ) -> Result<(), TransformError> {
245        if !done {
246            self.charge(text.len(), true)?;
247        }
248        let item = self.item(index, id)?;
249        let actual = match &item.value {
250            r::ResponseOutputItem::FunctionCall(_) => 0,
251            r::ResponseOutputItem::CustomToolCall(_) => 1,
252            r::ResponseOutputItem::CodeInterpreterCall(_) => 2,
253            r::ResponseOutputItem::McpCall(_) => 3,
254            _ => 255,
255        };
256        if actual != kind || item.argument_done {
257            return Err(invalid("argument event type/phase mismatch"));
258        }
259        let mut completed_bytes = 0;
260        if done {
261            if item.argument_streamed {
262                if item.arguments()? != text {
263                    return Err(invalid("argument done contradicts deltas"));
264                }
265            } else {
266                let suffix = text
267                    .strip_prefix(item.arguments()?)
268                    .ok_or_else(|| invalid("argument done contradicts initial arguments"))?;
269                completed_bytes = suffix.len();
270                item.arguments_mut()?.push_str(suffix);
271            }
272            item.argument_done = true;
273        } else {
274            item.arguments_mut()?.push_str(text);
275        }
276        item.argument_streamed = true;
277        self.charge(completed_bytes, true)
278    }
279    pub(super) fn charge(&mut self, n: usize, json: bool) -> Result<(), TransformError> {
280        let (used, cap) = if json {
281            (&mut self.json_bytes, self.limits.max_json_bytes)
282        } else {
283            (&mut self.text_bytes, self.limits.max_text_bytes)
284        };
285        *used = used.checked_add(n).ok_or_else(limit)?;
286        if *used > cap { Err(limit()) } else { Ok(()) }
287    }
288    pub(super) fn charge_item(
289        &mut self,
290        item: &r::ResponseOutputItem,
291    ) -> Result<(), TransformError> {
292        let (text, json) = item_sizes(item);
293        if text > self.limits.max_text_bytes || json > self.limits.max_json_bytes {
294            Err(limit())
295        } else {
296            Ok(())
297        }
298    }
299
300    pub(super) fn check_retained_sizes(&self) -> Result<(), TransformError> {
301        let (mut text, mut json) = (0usize, 0usize);
302        for item in &self.items {
303            let (t, j) = item_sizes(&item.value);
304            text = text.saturating_add(t);
305            json = json.saturating_add(j);
306            if !item.done {
307                for part in item.parts.values() {
308                    text = text.saturating_add(part.text().len());
309                }
310            }
311        }
312        if text > self.limits.max_text_bytes || json > self.limits.max_json_bytes {
313            Err(limit())
314        } else {
315            Ok(())
316        }
317    }
318    pub fn finish(self) -> Result<Converted<r::GenerateContentResponseBody>, TransformError> {
319        if self.failed || !self.done {
320            return Err(invalid("stream failed or ended before terminal response"));
321        }
322        Ok(Converted {
323            value: self.response.ok_or_else(|| invalid("missing response"))?,
324            report: self.report,
325        })
326    }
327}
328
329pub(super) fn item_sizes(item: &r::ResponseOutputItem) -> (usize, usize) {
330    let (mut text, mut json) = (0usize, 0usize);
331    match item {
332        r::ResponseOutputItem::Message(v) => {
333            for p in &v.content {
334                text = text.saturating_add(match p {
335                    i::OutputContent::Text(v) => v.text.len(),
336                    i::OutputContent::Refusal(v) => v.refusal.len(),
337                });
338            }
339        }
340        r::ResponseOutputItem::Reasoning(v) => {
341            for p in &v.summary {
342                text = text.saturating_add(p.text.len());
343            }
344            for p in v.content.iter().flatten() {
345                text = text.saturating_add(p.text.len());
346            }
347        }
348        r::ResponseOutputItem::FunctionCall(v) => json = v.arguments.len(),
349        r::ResponseOutputItem::CustomToolCall(v) => json = v.input.len(),
350        r::ResponseOutputItem::McpCall(v) => json = v.arguments.len(),
351        r::ResponseOutputItem::CodeInterpreterCall(v) => {
352            json = v.code.as_ref().map_or(0, String::len)
353        }
354        _ => {}
355    }
356    (text, json)
357}
gproxy_protocol/transform/generate/stream/responses/lifecycle.rs

gproxy_protocol/transform/generate/stream/responses/
lifecycle.rs
