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