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