1use super::{
2 items::{Item, item_id},
3 *,
4};
5use crate::{transform::Report, wire::DeclaredFields};
6use std::collections::HashSet;
7
8#[derive(Debug, Clone, Copy)]
9pub struct ResponsesStreamLimits {
10 pub max_bytes: usize,
11
12 pub max_text_bytes: usize,
13 pub max_json_bytes: usize,
14}
15
16impl Default for ResponsesStreamLimits {
17 fn default() -> Self {
18 Self {
19 max_bytes: 16 * 1024 * 1024,
20
21 max_text_bytes: 16 * 1024 * 1024,
22 max_json_bytes: 16 * 1024 * 1024,
23 }
24 }
25}
26
27pub struct ResponsesStreamCollector {
28 pub(super) limits: ResponsesStreamLimits,
29 pub(super) sequence: Option<i64>,
30 pub(super) bytes: usize,
31
32 pub(super) items: Vec<Item>,
33 pub(super) ids: HashSet<String>,
34 pub(super) response: Option<r::GenerateContentResponseBody>,
35 pub(super) done: bool,
36 pub(super) failed: bool,
37 pub(super) text_bytes: usize,
38 pub(super) json_bytes: usize,
39 pub(super) report: Report,
40}
41
42impl ResponsesStreamCollector {
43 pub fn new(limits: ResponsesStreamLimits) -> Self {
44 Self {
45 limits,
46 sequence: None,
47 bytes: 0,
48
49 items: Vec::new(),
50 ids: HashSet::new(),
51 response: None,
52 done: false,
53 failed: false,
54 text_bytes: 0,
55 json_bytes: 0,
56 report: Report::default(),
57 }
58 }
59 pub fn push(&mut self, event: s::StreamEvent) -> Result<(), TransformError> {
60 if self.done || self.failed {
61 self.failed = true;
62 return Err(invalid("event after terminal/failed stream"));
63 }
64 let result = self.push_inner(event.into_declared());
65 if result.is_err() {
66 self.failed = true;
67 }
68 result
69 }
70 fn push_inner(&mut self, event: s::StreamEvent) -> Result<(), TransformError> {
71 let event = crate::transform::generate::multi_agent::attribute_event(event)?;
72 if let (Some(agent), Some(index)) = (event.agent(), event.output_index())
73 && let Some(item) = usize::try_from(index)
74 .ok()
75 .and_then(|index| self.items.get(index))
76 && item
77 .value
78 .agent()
79 .map_or("/root", |owner| owner.agent_name.as_str())
80 != agent.agent_name
81 {
82 return Err(invalid("event agent differs from its output item"));
83 }
84
85 self.bytes += bounded(&event, self.limits.max_bytes.saturating_sub(self.bytes))?;
86 let Some(seq) = super::events::sequence(&event) else {
87 return Ok(());
88 };
89 if seq < 0 || self.sequence.is_some_and(|v| seq <= v) {
90 return Err(invalid("nonmonotonic sequence number"));
91 }
92 self.sequence = Some(seq);
93 if self.response.is_none()
94 && !matches!(event, s::StreamEvent::Created(_) | s::StreamEvent::Error(_))
95 {
96 return Err(invalid("response.created required before other events"));
97 }
98 match event {
99 s::StreamEvent::Keepalive => {}
100 s::StreamEvent::Created(v) => {
101 if self.response.is_some() {
102 return Err(invalid("duplicate/nonempty response.created"));
103 }
104 self.check_response(&v.response, false)?;
105 self.seed_output(&v.response.output)?;
106 self.response = Some(v.response);
107 }
108 s::StreamEvent::Queued(v) => {
109 self.check_response(&v.response, false)?;
110 self.snapshot_output(&v.response.output)?;
111 if v.response.status != Some(r::ResponseStatus::Queued) {
112 return Err(invalid("queued status mismatch"));
113 }
114 }
115 s::StreamEvent::InProgress(v) => {
116 self.check_response(&v.response, false)?;
117 self.snapshot_output(&v.response.output)?;
118 if v.response.status != Some(r::ResponseStatus::InProgress) {
119 return Err(invalid("in_progress status mismatch"));
120 }
121 }
122 s::StreamEvent::Completed(v) => {
123 self.terminal(v.response, r::ResponseStatus::Completed)?
124 }
125 s::StreamEvent::Incomplete(v) => {
126 self.terminal(v.response, r::ResponseStatus::Incomplete)?
127 }
128 s::StreamEvent::Failed(value) => {
129 self.check_response(&value.response, true)?;
130 if value.response.status != Some(r::ResponseStatus::Failed) {
131 return Err(invalid("failed event status mismatch"));
132 }
133 let detail = value
134 .response
135 .error
136 .map(|error| format!("{:?}: {}", error.code, error.message))
137 .unwrap_or_else(|| "upstream response failed without error detail".into());
138 return Err(TransformError::invalid_result(
139 "responses.stream.failed",
140 detail,
141 ));
142 }
143 s::StreamEvent::Error(error) => {
144 return Err(TransformError::invalid_result(
145 "responses.stream.error",
146 format!(
147 "{}{}: {}",
148 error.code.as_deref().unwrap_or("error"),
149 error
150 .param
151 .as_ref()
152 .map(|p| format!(" ({p})"))
153 .unwrap_or_default(),
154 error.message
155 ),
156 ));
157 }
158 s::StreamEvent::OutputItemAdded(v) => {
159 if v.output_index != self.items.len() as i64 {
160 return Err(invalid("noncontiguous/duplicate output index"));
161 }
162
163 if let Some(id) = item_id(&v.item)
164 && (id.is_empty() || !self.ids.insert(id.to_owned()))
165 {
166 return Err(invalid("empty/duplicate item id"));
167 }
168 self.charge_item(&v.item)?;
169 let (text, json) = super::lifecycle::item_sizes(&v.item);
170 self.charge(text, false)?;
171 self.charge(json, true)?;
172 self.items.push(Item::new(v.item));
173 }
174 s::StreamEvent::OutputItemDone(v) => {
175 self.charge_item(&v.item)?;
176 if let Some(id) = item_id(&v.item) {
177 self.item(v.output_index, id)?;
178 }
179 let item = self
180 .items
181 .get_mut(
182 usize::try_from(v.output_index)
183 .map_err(|_| invalid("negative output index"))?,
184 )
185 .ok_or_else(|| invalid("item done without added"))?;
186 item.finish(v.item)?;
187 self.check_retained_sizes()?;
188 }
189 s::StreamEvent::ContentPartAdded(v) => {
190 self.add_part(v.output_index, &v.item_id, v.content_index, false, v.part)?
191 }
192 s::StreamEvent::ContentPartDone(v) => self
193 .part(v.output_index, &v.item_id, v.content_index, false)?
194 .finish(v.part)?,
195 s::StreamEvent::ReasoningSummaryPartAdded(v) => self.add_part(
196 v.output_index,
197 &v.item_id,
198 v.summary_index,
199 true,
200 s::OutputContentPart::Reasoning(s::ReasoningText {
201 text: v.part.text,
202 type_: s::ReasoningTextType::ReasoningText,
203 rest: Default::default(),
204 }),
205 )?,
206 s::StreamEvent::ReasoningSummaryPartDone(v) => {
207 let part = self.part(v.output_index, &v.item_id, v.summary_index, true)?;
208 part.incomplete = v.status == Some(s::SummaryPartStatus::Incomplete);
209 part.finish(s::OutputContentPart::Reasoning(s::ReasoningText {
210 text: v.part.text,
211 type_: s::ReasoningTextType::ReasoningText,
212 rest: Default::default(),
213 }))?;
214 }
215 s::StreamEvent::OutputTextDelta(v) => {
216 self.delta(v.output_index, &v.item_id, v.content_index, 0, &v.delta)?
217 .logs
218 .extend(v.logprobs);
219 }
220 s::StreamEvent::RefusalDelta(v) => {
221 self.delta(v.output_index, &v.item_id, v.content_index, 1, &v.delta)?;
222 }
223 s::StreamEvent::ReasoningTextDelta(v) => {
224 self.delta(v.output_index, &v.item_id, v.content_index, 2, &v.delta)?;
225 }
226 s::StreamEvent::ReasoningSummaryTextDelta(v) => {
227 self.delta(v.output_index, &v.item_id, v.summary_index, 3, &v.delta)?;
228 }
229 s::StreamEvent::OutputTextDone(v) => self.text_done(
230 v.output_index,
231 &v.item_id,
232 v.content_index,
233 0,
234 &v.text,
235 &v.logprobs,
236 )?,
237 s::StreamEvent::RefusalDone(v) => self.text_done(
238 v.output_index,
239 &v.item_id,
240 v.content_index,
241 1,
242 &v.refusal,
243 &[],
244 )?,
245 s::StreamEvent::ReasoningTextDone(v) => {
246 self.text_done(v.output_index, &v.item_id, v.content_index, 2, &v.text, &[])?
247 }
248 s::StreamEvent::ReasoningSummaryTextDone(v) => {
249 self.text_done(v.output_index, &v.item_id, v.summary_index, 3, &v.text, &[])?
250 }
251 s::StreamEvent::OutputTextAnnotationAdded(v) => {
252 let p = self.part(v.output_index, &v.item_id, v.content_index, false)?;
253 if p.done {
254 return Err(invalid("annotation after part done"));
255 }
256 let s::OutputContentPart::Text(text) = &mut p.value else {
257 return Err(invalid("annotation on nontext part"));
258 };
259 if v.annotation_index != p.annotation_index as i64 {
260 return Err(invalid("noncontiguous annotation index"));
261 }
262 p.annotation_index += 1;
263 let known = v
264 .annotation
265 .get("type")
266 .and_then(|v| v.as_str())
267 .is_some_and(|v| {
268 matches!(
269 v,
270 "url_citation"
271 | "file_citation"
272 | "container_file_citation"
273 | "file_path"
274 )
275 });
276 match serde_json::from_value::<i::OutputAnnotation>(v.annotation) {
277 Ok(annotation) => text.annotations.push(annotation.into_declared()),
278 Err(_) if known => return Err(invalid("malformed known annotation")),
279 Err(_) => self.report.omitted(
280 "response.output_text.annotation.added",
281 "unknown annotation has no typed final output representation",
282 ),
283 }
284 }
285 s::StreamEvent::FunctionCallArgumentsDelta(v) => {
286 self.argument(v.output_index, &v.item_id, 0, &v.delta, false)?
287 }
288 s::StreamEvent::FunctionCallArgumentsDone(v) => {
289 if !matches!(&self.item(v.output_index,&v.item_id)?.value,r::ResponseOutputItem::FunctionCall(c) if c.name==v.name)
290 {
291 return Err(invalid("function done name mismatch"));
292 }
293 self.argument(v.output_index, &v.item_id, 0, &v.arguments, true)?;
294 }
295 s::StreamEvent::CustomToolInputDelta(v) => {
296 self.argument(v.output_index, &v.item_id, 1, &v.delta, false)?
297 }
298 s::StreamEvent::CustomToolInputDone(v) => {
299 self.argument(v.output_index, &v.item_id, 1, &v.input, true)?
300 }
301 s::StreamEvent::CodeInterpreterCodeDelta(v) => {
302 self.argument(v.output_index, &v.item_id, 2, &v.delta, false)?
303 }
304 s::StreamEvent::CodeInterpreterCodeDone(v) => {
305 self.argument(v.output_index, &v.item_id, 2, &v.code, true)?
306 }
307 s::StreamEvent::McpArgumentsDelta(v) => {
308 self.argument(v.output_index, &v.item_id, 3, &v.delta, false)?
309 }
310 s::StreamEvent::McpArgumentsDone(v) => {
311 self.argument(v.output_index, &v.item_id, 3, &v.arguments, true)?
312 }
313 s::StreamEvent::ImageCall(v) => {
314 self.item(v.output_index, &v.item_id)?.progress(0, 0)?
315 }
316 s::StreamEvent::ImageGenerating(v) => {
317 self.item(v.output_index, &v.item_id)?.progress(0, 1)?
318 }
319 s::StreamEvent::ImageCompleted(v) => {
320 self.item(v.output_index, &v.item_id)?.progress(0, 2)?
321 }
322 s::StreamEvent::CodeInterpreterInProgress(v) => {
323 self.item(v.output_index, &v.item_id)?.progress(1, 0)?
324 }
325 s::StreamEvent::CodeInterpreterInterpreting(v) => {
326 self.item(v.output_index, &v.item_id)?.progress(1, 1)?
327 }
328 s::StreamEvent::CodeInterpreterCompleted(v) => {
329 self.item(v.output_index, &v.item_id)?.progress(1, 2)?
330 }
331 s::StreamEvent::FileSearchInProgress(v) => {
332 self.item(v.output_index, &v.item_id)?.progress(2, 0)?
333 }
334 s::StreamEvent::FileSearchSearching(v) => {
335 self.item(v.output_index, &v.item_id)?.progress(2, 1)?
336 }
337 s::StreamEvent::FileSearchCompleted(v) => {
338 self.item(v.output_index, &v.item_id)?.progress(2, 2)?
339 }
340 s::StreamEvent::WebSearchInProgress(v) => {
341 self.item(v.output_index, &v.item_id)?.progress(3, 0)?
342 }
343 s::StreamEvent::WebSearchSearching(v) => {
344 self.item(v.output_index, &v.item_id)?.progress(3, 1)?
345 }
346 s::StreamEvent::WebSearchCompleted(v) => {
347 self.item(v.output_index, &v.item_id)?.progress(3, 2)?
348 }
349 s::StreamEvent::McpInProgress(v) => {
350 self.item(v.output_index, &v.item_id)?.progress(4, 0)?
351 }
352 s::StreamEvent::McpCompleted(v) => {
353 self.item(v.output_index, &v.item_id)?.progress(4, 2)?
354 }
355 s::StreamEvent::McpFailed(v) => {
356 self.item(v.output_index, &v.item_id)?.progress(4, 3)?
357 }
358 s::StreamEvent::McpListToolsInProgress(v) => {
359 self.item(v.output_index, &v.item_id)?.progress(5, 0)?
360 }
361 s::StreamEvent::McpListToolsCompleted(v) => {
362 self.item(v.output_index, &v.item_id)?.progress(5, 2)?
363 }
364 s::StreamEvent::McpListToolsFailed(v) => {
365 self.item(v.output_index, &v.item_id)?.progress(5, 3)?
366 }
367 s::StreamEvent::ImagePartial(v) => {
368 if v.partial_image_index < 0
369 || !matches!(
370 self.item(v.output_index, &v.item_id)?.value,
371 r::ResponseOutputItem::ImageGenerationCall(_)
372 )
373 {
374 return Err(invalid("image preview identity/type mismatch"));
375 }
376 if !self
377 .report
378 .diagnostics
379 .iter()
380 .any(|d| d.field == "response.image_generation_call.partial_image")
381 {
382 self.report.omitted(
383 "response.image_generation_call.partial_image",
384 "preview frames are superseded by the canonical image result",
385 );
386 }
387 }
388 s::StreamEvent::AudioDelta(_)
389 | s::StreamEvent::AudioDone(_)
390 | s::StreamEvent::AudioTranscriptDelta(_)
391 | s::StreamEvent::AudioTranscriptDone(_) => {
392 return Err(TransformError::unsupported(
393 "response.audio",
394 "native response output has no corresponding audio item; use audio event consumer",
395 ));
396 }
397 }
398 Ok(())
399 }
400}