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