1use std::borrow::Cow;
33use std::collections::VecDeque;
34use std::fmt;
35use std::pin::Pin;
36use std::task::{Context, Poll};
37
38use futures_core::Stream;
39use serde::{Deserialize, Serialize};
40
41use crate::chat::content::is_valid_json;
42pub use crate::chat::transcript::StreamTranscript;
43use crate::{
44 ChatResponse, Citation, ContentBlock, Error, ExtraMap, FinishReason, Image, MediaSource,
45 OtherContent, ProviderData, ProviderToolCall, ProviderToolResult, Reasoning, Refusal, Result,
46 StreamErrorKind, Text, ToolCall, Usage, Warning,
47};
48
49pub type ChatStream = Pin<Box<dyn Stream<Item = Result<StreamEvent>> + Send + 'static>>;
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
70#[serde(rename_all = "snake_case")]
71pub enum UsageMetadataMode {
72 Delta,
74 #[default]
76 Snapshot,
77}
78
79fn is_snapshot(mode: &UsageMetadataMode) -> bool {
80 *mode == UsageMetadataMode::Snapshot
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86#[non_exhaustive]
87pub enum StreamBlockType {
88 Text,
90 Image,
92 Reasoning,
94 Refusal,
96 ToolCall,
100 ProviderToolCall,
105 ProviderToolResult,
108 Other,
110}
111
112impl StreamBlockType {
113 pub const ALL: &'static [Self] = &[
116 Self::Text,
117 Self::Image,
118 Self::Reasoning,
119 Self::Refusal,
120 Self::ToolCall,
121 Self::ProviderToolCall,
122 Self::ProviderToolResult,
123 Self::Other,
124 ];
125
126 fn label(self) -> &'static str {
128 match self {
129 Self::Text => "text",
130 Self::Image => "image",
131 Self::Reasoning => "reasoning",
132 Self::Refusal => "refusal",
133 Self::ToolCall => "tool_call",
134 Self::ProviderToolCall => "provider_tool_call",
135 Self::ProviderToolResult => "provider_tool_result",
136 Self::Other => "other",
137 }
138 }
139}
140
141#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
201#[serde(tag = "event", rename_all = "snake_case")]
202#[non_exhaustive]
203pub enum StreamEvent {
204 #[non_exhaustive]
206 ResponseStart {
207 #[serde(default, skip_serializing_if = "Option::is_none")]
209 id: Option<String>,
210 #[serde(default, skip_serializing_if = "Option::is_none")]
212 model: Option<String>,
213 },
214 #[doc(alias = "ToolCallStart")]
217 #[non_exhaustive]
218 BlockStart {
219 index: usize,
221 block_type: StreamBlockType,
223 #[serde(default, skip_serializing_if = "Option::is_none")]
225 id: Option<String>,
226 #[serde(default, skip_serializing_if = "Option::is_none")]
228 name: Option<String>,
229 #[serde(default, skip_serializing_if = "Option::is_none")]
231 type_name: Option<String>,
232 #[serde(default, skip_serializing_if = "Option::is_none")]
234 data: Option<ExtraMap>,
235 #[serde(default, skip_serializing_if = "Option::is_none")]
237 origin: Option<String>,
238 #[serde(default, skip_serializing_if = "ProviderData::is_empty")]
240 provider_data: ProviderData,
241 },
242 #[non_exhaustive]
244 TextDelta {
245 index: usize,
247 text: String,
249 },
250 #[non_exhaustive]
252 ReasoningDelta {
253 index: usize,
255 text: String,
257 },
258 #[non_exhaustive]
265 RefusalDelta {
266 index: usize,
268 text: String,
270 },
271 #[non_exhaustive]
273 ToolCallDelta {
274 index: usize,
276 arguments: String,
278 },
279 #[non_exhaustive]
282 Citation {
283 index: usize,
285 citation: Citation,
287 },
288 #[non_exhaustive]
290 BlockProviderData {
291 index: usize,
293 provider_data: ProviderData,
295 },
296 #[non_exhaustive]
305 ResponseProviderData {
306 provider_data: ProviderData,
308 },
309 #[non_exhaustive]
318 ProviderDelta {
319 index: usize,
321 type_name: String,
323 #[serde(default, skip_serializing_if = "ExtraMap::is_empty")]
325 data: ExtraMap,
326 #[serde(default, skip_serializing_if = "Option::is_none")]
328 origin: Option<String>,
329 },
330 #[non_exhaustive]
332 BlockStop {
333 index: usize,
335 },
336 #[non_exhaustive]
338 Warning {
339 warning: Warning,
341 },
342 #[non_exhaustive]
344 ResponseMetadata {
345 #[serde(default, skip_serializing_if = "Option::is_none")]
347 finish_reason: Option<FinishReason>,
348 #[serde(default, skip_serializing_if = "Option::is_none")]
350 finish_reason_raw: Option<String>,
351 #[serde(default, skip_serializing_if = "Option::is_none")]
353 usage: Option<Usage>,
354 #[serde(default, skip_serializing_if = "is_snapshot")]
356 usage_mode: UsageMetadataMode,
357 #[serde(default, skip_serializing_if = "Option::is_none")]
359 id: Option<String>,
360 #[serde(default, skip_serializing_if = "Option::is_none")]
362 model: Option<String>,
363 #[serde(default, skip_serializing_if = "Option::is_none")]
365 request_id: Option<String>,
366 #[serde(default, skip_serializing_if = "Option::is_none")]
368 provider: Option<Cow<'static, str>>,
369 #[serde(default, skip_serializing_if = "ProviderData::is_empty")]
371 provider_data: ProviderData,
372 },
373 ResponseStop,
379}
380
381impl StreamEvent {
382 #[must_use]
384 pub fn response_start(id: Option<String>, model: Option<String>) -> Self {
385 Self::ResponseStart { id, model }
386 }
387
388 #[must_use]
390 pub fn text_delta(index: usize, text: impl Into<String>) -> Self {
391 Self::TextDelta {
392 index,
393 text: text.into(),
394 }
395 }
396
397 #[must_use]
399 pub fn reasoning_delta(index: usize, text: impl Into<String>) -> Self {
400 Self::ReasoningDelta {
401 index,
402 text: text.into(),
403 }
404 }
405
406 #[must_use]
408 pub fn refusal_delta(index: usize, text: impl Into<String>) -> Self {
409 Self::RefusalDelta {
410 index,
411 text: text.into(),
412 }
413 }
414
415 #[must_use]
417 pub fn tool_call_delta(index: usize, arguments: impl Into<String>) -> Self {
418 Self::ToolCallDelta {
419 index,
420 arguments: arguments.into(),
421 }
422 }
423
424 #[must_use]
426 pub fn citation(index: usize, citation: Citation) -> Self {
427 Self::Citation { index, citation }
428 }
429
430 #[must_use]
432 pub fn block_provider_data(index: usize, provider_data: ProviderData) -> Self {
433 Self::BlockProviderData {
434 index,
435 provider_data,
436 }
437 }
438
439 #[must_use]
442 pub fn response_provider_data(provider_data: ProviderData) -> Self {
443 Self::ResponseProviderData { provider_data }
444 }
445
446 #[must_use]
448 pub fn block_stop(index: usize) -> Self {
449 Self::BlockStop { index }
450 }
451
452 #[must_use]
454 pub fn warning(warning: Warning) -> Self {
455 Self::Warning { warning }
456 }
457
458 #[must_use]
460 pub fn metadata() -> Self {
461 Self::ResponseMetadata {
462 finish_reason: None,
463 finish_reason_raw: None,
464 usage: None,
465 usage_mode: UsageMetadataMode::Snapshot,
466 id: None,
467 model: None,
468 request_id: None,
469 provider: None,
470 provider_data: ProviderData::new(),
471 }
472 }
473
474 #[must_use]
486 pub fn block_start(index: usize, block_type: StreamBlockType) -> Self {
487 Self::BlockStart {
488 index,
489 block_type,
490 id: None,
491 name: None,
492 type_name: None,
493 data: None,
494 origin: None,
495 provider_data: ProviderData::new(),
496 }
497 }
498
499 #[must_use]
502 pub fn provider_tool_call_start(
503 index: usize,
504 id: impl Into<String>,
505 name: impl Into<String>,
506 origin: Option<String>,
507 ) -> Self {
508 Self::BlockStart {
509 index,
510 block_type: StreamBlockType::ProviderToolCall,
511 id: Some(id.into()),
512 name: Some(name.into()),
513 type_name: None,
514 data: None,
515 origin,
516 provider_data: ProviderData::new(),
517 }
518 }
519
520 #[must_use]
523 pub fn provider_tool_result_start(index: usize, result: ProviderToolResult) -> Self {
524 let mut data = ExtraMap::new();
525 data.insert("output".into(), result.output);
526 if result.is_error {
527 data.insert("is_error".into(), true.into());
528 }
529 Self::BlockStart {
530 index,
531 block_type: StreamBlockType::ProviderToolResult,
532 id: Some(result.call_id),
533 name: None,
534 type_name: None,
535 data: Some(data),
536 origin: result.origin,
537 provider_data: result.provider_data,
538 }
539 }
540
541 #[must_use]
543 pub fn image_start(index: usize, image: Image) -> Self {
544 Self::BlockStart {
545 index,
546 block_type: StreamBlockType::Image,
547 id: None,
548 name: None,
549 type_name: None,
550 data: Some(image_data(&image.source, image.media_type.as_deref())),
551 origin: None,
552 provider_data: image.provider_data,
553 }
554 }
555
556 #[must_use]
559 pub fn other_start(index: usize, other: OtherContent) -> Self {
560 Self::BlockStart {
561 index,
562 block_type: StreamBlockType::Other,
563 id: None,
564 name: None,
565 type_name: Some(other.type_name),
566 data: Some(other.data),
567 origin: other.origin,
568 provider_data: other.provider_data,
569 }
570 }
571
572 #[must_use]
574 pub fn provider_delta(
575 index: usize,
576 type_name: impl Into<String>,
577 data: ExtraMap,
578 origin: Option<String>,
579 ) -> Self {
580 Self::ProviderDelta {
581 index,
582 type_name: type_name.into(),
583 data,
584 origin,
585 }
586 }
587
588 #[must_use]
590 pub fn tool_call_start(index: usize, id: impl Into<String>, name: impl Into<String>) -> Self {
591 Self::BlockStart {
592 index,
593 block_type: StreamBlockType::ToolCall,
594 id: Some(id.into()),
595 name: Some(name.into()),
596 type_name: None,
597 data: None,
598 origin: None,
599 provider_data: ProviderData::new(),
600 }
601 }
602
603 #[must_use]
609 pub fn is_output(&self) -> bool {
610 match self {
611 Self::TextDelta { .. }
612 | Self::ReasoningDelta { .. }
613 | Self::RefusalDelta { .. }
614 | Self::ToolCallDelta { .. } => true,
615 Self::BlockStart { block_type, .. } => {
616 matches!(
617 block_type,
618 StreamBlockType::Image
619 | StreamBlockType::Refusal
620 | StreamBlockType::ProviderToolResult
621 | StreamBlockType::Other
622 )
623 }
624 _ => false,
625 }
626 }
627
628 fn is_block_event(&self) -> bool {
630 matches!(
631 self,
632 Self::BlockStart { .. }
633 | Self::TextDelta { .. }
634 | Self::ReasoningDelta { .. }
635 | Self::RefusalDelta { .. }
636 | Self::ToolCallDelta { .. }
637 | Self::Citation { .. }
638 | Self::BlockProviderData { .. }
639 | Self::ProviderDelta { .. }
640 | Self::BlockStop { .. }
641 )
642 }
643
644 pub(crate) fn name(&self) -> &'static str {
645 match self {
646 Self::ResponseStart { .. } => "response_start",
647 Self::BlockStart { .. } => "block_start",
648 Self::TextDelta { .. } => "text_delta",
649 Self::ReasoningDelta { .. } => "reasoning_delta",
650 Self::RefusalDelta { .. } => "refusal_delta",
651 Self::ToolCallDelta { .. } => "tool_call_delta",
652 Self::Citation { .. } => "citation",
653 Self::BlockProviderData { .. } => "block_provider_data",
654 Self::ResponseProviderData { .. } => "response_provider_data",
655 Self::ProviderDelta { .. } => "provider_delta",
656 Self::BlockStop { .. } => "block_stop",
657 Self::Warning { .. } => "warning",
658 Self::ResponseMetadata { .. } => "response_metadata",
659 Self::ResponseStop => "response_stop",
660 }
661 }
662}
663
664#[derive(Debug, Clone, PartialEq, Eq)]
709#[must_use]
710pub struct MetadataBuilder {
711 event: StreamEvent,
713}
714
715impl Default for MetadataBuilder {
716 fn default() -> Self {
717 Self::new()
718 }
719}
720
721impl MetadataBuilder {
722 pub fn new() -> Self {
724 Self {
725 event: StreamEvent::metadata(),
726 }
727 }
728
729 pub fn with_finish_reason(self, value: FinishReason) -> Self {
731 self.with_finish_reason_opt(Some(value))
732 }
733
734 pub fn with_finish_reason_opt(mut self, value: Option<FinishReason>) -> Self {
738 if let StreamEvent::ResponseMetadata { finish_reason, .. } = &mut self.event {
739 merge(finish_reason, value);
740 }
741 self
742 }
743
744 pub fn with_finish_reason_raw(mut self, value: impl Into<String>) -> Self {
746 if let StreamEvent::ResponseMetadata {
747 finish_reason_raw, ..
748 } = &mut self.event
749 {
750 *finish_reason_raw = Some(value.into());
751 }
752 self
753 }
754
755 pub fn with_finish_reason_raw_opt(self, value: Option<&str>) -> Self {
759 match value {
760 Some(value) => self.with_finish_reason_raw(value),
761 None => self,
762 }
763 }
764
765 pub fn with_usage(self, value: Usage) -> Self {
767 self.with_usage_opt(Some(value))
768 }
769
770 pub fn with_usage_opt(mut self, value: Option<Usage>) -> Self {
773 if let StreamEvent::ResponseMetadata { usage, .. } = &mut self.event {
774 merge(usage, value);
775 }
776 self
777 }
778
779 pub fn with_usage_mode(mut self, value: UsageMetadataMode) -> Self {
781 if let StreamEvent::ResponseMetadata { usage_mode, .. } = &mut self.event {
782 *usage_mode = value;
783 }
784 self
785 }
786
787 pub fn with_id(mut self, value: impl Into<String>) -> Self {
789 if let StreamEvent::ResponseMetadata { id, .. } = &mut self.event {
790 *id = Some(value.into());
791 }
792 self
793 }
794
795 pub fn with_id_opt(self, value: Option<&str>) -> Self {
798 match value {
799 Some(value) => self.with_id(value),
800 None => self,
801 }
802 }
803
804 pub fn with_model(mut self, value: impl Into<String>) -> Self {
806 if let StreamEvent::ResponseMetadata { model, .. } = &mut self.event {
807 *model = Some(value.into());
808 }
809 self
810 }
811
812 pub fn with_model_opt(self, value: Option<&str>) -> Self {
815 match value {
816 Some(value) => self.with_model(value),
817 None => self,
818 }
819 }
820
821 pub fn with_request_id(mut self, value: impl Into<String>) -> Self {
823 if let StreamEvent::ResponseMetadata { request_id, .. } = &mut self.event {
824 *request_id = Some(value.into());
825 }
826 self
827 }
828
829 pub fn with_request_id_opt(self, value: Option<&str>) -> Self {
833 match value {
834 Some(value) => self.with_request_id(value),
835 None => self,
836 }
837 }
838
839 pub fn with_provider(self, value: impl Into<Cow<'static, str>>) -> Self {
846 self.with_provider_opt(Some(value.into()))
847 }
848
849 pub fn with_provider_opt(mut self, value: Option<Cow<'static, str>>) -> Self {
854 if let StreamEvent::ResponseMetadata { provider, .. } = &mut self.event {
855 merge(provider, value);
856 }
857 self
858 }
859
860 pub fn with_provider_data(mut self, value: ProviderData) -> Self {
862 if let StreamEvent::ResponseMetadata { provider_data, .. } = &mut self.event {
863 *provider_data = value;
864 }
865 self
866 }
867
868 #[must_use]
870 pub fn build(self) -> StreamEvent {
871 self.event
872 }
873}
874
875impl From<MetadataBuilder> for StreamEvent {
876 fn from(builder: MetadataBuilder) -> Self {
877 builder.build()
878 }
879}
880
881#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
883#[non_exhaustive]
884pub struct StreamCompleteness {
885 pub saw_response_start: bool,
887 pub saw_response_stop: bool,
889 pub has_open_blocks: bool,
891 pub dropped_incomplete_tool_calls: bool,
893}
894
895impl StreamCompleteness {
896 #[must_use]
898 pub fn is_complete(&self) -> bool {
899 self.saw_response_start
900 && self.saw_response_stop
901 && !self.has_open_blocks
902 && !self.dropped_incomplete_tool_calls
903 }
904}
905
906#[derive(Debug, Clone)]
908#[non_exhaustive]
909pub struct CollectedResponse {
910 pub response: ChatResponse,
912 pub completeness: StreamCompleteness,
914 pub terminal_error: Option<Error>,
916}
917
918#[derive(Default)]
935pub struct StreamCollector {
936 blocks: Vec<Option<Slot>>,
937 skipped_indices: usize,
939 response: ChatResponse,
940 saw_response_start: bool,
941 saw_response_stop: bool,
942 terminal_error: Option<Error>,
943 metadata_only: bool,
945}
946
947enum Slot {
948 Open(BlockAccumulator),
949 Closed {
950 block_type: StreamBlockType,
951 block: Finished,
952 },
953}
954
955enum Finished {
957 Block(ContentBlock),
958 Empty,
960 InvalidToolCall(ToolCall),
962 IncompleteToolCall,
964}
965
966struct BlockAccumulator {
967 kind: BlockKind,
968 provider_data: ProviderData,
969 started: bool,
970}
971
972impl BlockAccumulator {
973 fn finish(self, closed: bool) -> Finished {
974 let provider_data = self.provider_data;
975 match self.kind {
976 BlockKind::Text { text, citations }
977 if text.is_empty() && citations.is_empty() && provider_data.is_empty() =>
978 {
979 Finished::Empty
980 }
981 BlockKind::Text { text, citations } => Finished::Block(ContentBlock::Text(
982 Text::new(text)
983 .with_citations(citations)
984 .with_provider_data(provider_data),
985 )),
986 BlockKind::Image { source, media_type } => {
987 let mut image = Image::new(source);
988 image.media_type = media_type;
989 image.provider_data = provider_data;
990 Finished::Block(ContentBlock::Image(image))
991 }
992 BlockKind::Reasoning(text) if text.is_empty() && provider_data.is_empty() => {
993 Finished::Empty
994 }
995 BlockKind::Reasoning(text) => Finished::Block(ContentBlock::Reasoning(
996 Reasoning::new(text).with_provider_data(provider_data),
997 )),
998 BlockKind::Refusal(text) if text.is_empty() && provider_data.is_empty() => {
999 Finished::Empty
1000 }
1001 BlockKind::Refusal(text) => Finished::Block(ContentBlock::Refusal(
1002 Refusal::new(text).with_provider_data(provider_data),
1003 )),
1004 BlockKind::ToolCall {
1005 id: Some(id),
1006 name: Some(name),
1007 arguments,
1008 } if closed => {
1009 let arguments = if arguments.trim().is_empty() {
1010 "{}".to_owned()
1011 } else {
1012 arguments
1013 };
1014 let valid = is_valid_json(&arguments);
1015 let call = ToolCall::new(id, name, arguments).with_provider_data(provider_data);
1016 if valid {
1017 Finished::Block(ContentBlock::ToolCall(call))
1018 } else {
1019 Finished::InvalidToolCall(call)
1020 }
1021 }
1022 BlockKind::ToolCall { .. } => Finished::IncompleteToolCall,
1023 BlockKind::ProviderToolCall {
1024 id: Some(id),
1025 name: Some(name),
1026 arguments,
1027 origin,
1028 } if closed => {
1029 let arguments = if arguments.trim().is_empty() {
1030 "{}".to_owned()
1031 } else {
1032 arguments
1033 };
1034 let mut call =
1035 ProviderToolCall::new(id, name, arguments).with_provider_data(provider_data);
1036 call.origin = origin;
1037 Finished::Block(ContentBlock::ProviderToolCall(call))
1038 }
1039 BlockKind::ProviderToolCall { .. } => Finished::IncompleteToolCall,
1040 BlockKind::ProviderToolResult {
1041 call_id,
1042 output,
1043 is_error,
1044 origin,
1045 } => {
1046 let mut result = ProviderToolResult::new(call_id, output)
1047 .with_error(is_error)
1048 .with_provider_data(provider_data);
1049 result.origin = origin;
1050 Finished::Block(ContentBlock::ProviderToolResult(result))
1051 }
1052 BlockKind::Other {
1053 type_name,
1054 data,
1055 origin,
1056 } => {
1057 let mut other =
1058 OtherContent::new(type_name, data).with_provider_data(provider_data);
1059 other.origin = origin;
1060 Finished::Block(ContentBlock::Other(other))
1061 }
1062 }
1063 }
1064}
1065
1066enum BlockKind {
1067 Text {
1068 text: String,
1069 citations: Vec<Citation>,
1070 },
1071 Image {
1072 source: MediaSource,
1073 media_type: Option<String>,
1075 },
1076 Reasoning(String),
1077 Refusal(String),
1078 ToolCall {
1079 id: Option<String>,
1080 name: Option<String>,
1081 arguments: String,
1082 },
1083 ProviderToolCall {
1084 id: Option<String>,
1085 name: Option<String>,
1086 arguments: String,
1087 origin: Option<String>,
1088 },
1089 ProviderToolResult {
1090 call_id: String,
1091 output: serde_json::Value,
1092 is_error: bool,
1093 origin: Option<String>,
1094 },
1095 Other {
1096 type_name: String,
1097 data: ExtraMap,
1098 origin: Option<String>,
1099 },
1100}
1101
1102impl BlockKind {
1103 fn block_type(&self) -> StreamBlockType {
1104 match self {
1105 Self::Text { .. } => StreamBlockType::Text,
1106 Self::Image { .. } => StreamBlockType::Image,
1107 Self::Reasoning(_) => StreamBlockType::Reasoning,
1108 Self::Refusal(_) => StreamBlockType::Refusal,
1109 Self::ToolCall { .. } => StreamBlockType::ToolCall,
1110 Self::ProviderToolCall { .. } => StreamBlockType::ProviderToolCall,
1111 Self::ProviderToolResult { .. } => StreamBlockType::ProviderToolResult,
1112 Self::Other { .. } => StreamBlockType::Other,
1113 }
1114 }
1115
1116 fn label(&self) -> &'static str {
1117 self.block_type().label()
1118 }
1119
1120 fn text() -> Self {
1121 Self::Text {
1122 text: String::new(),
1123 citations: Vec::new(),
1124 }
1125 }
1126}
1127
1128impl StreamCollector {
1129 #[must_use]
1131 pub fn new() -> Self {
1132 Self::default()
1133 }
1134
1135 #[must_use]
1180 pub fn metadata_only() -> Self {
1181 Self {
1182 metadata_only: true,
1183 ..Self::default()
1184 }
1185 }
1186
1187 pub fn push(&mut self, event: StreamEvent) -> Result<Option<&ContentBlock>> {
1194 if self.saw_response_stop {
1195 return Err(Error::stream_protocol(match event {
1196 StreamEvent::ResponseStop => "duplicate response_stop".to_owned(),
1197 other => format!("{} received after response_stop", other.name()),
1198 }));
1199 }
1200
1201 if self.metadata_only && event.is_block_event() {
1202 return Ok(None);
1203 }
1204
1205 match event {
1206 StreamEvent::ResponseStart { id, model } => {
1207 self.saw_response_start = true;
1208 merge(&mut self.response.id, id);
1209 merge(&mut self.response.model, model);
1210 }
1211 StreamEvent::BlockStart {
1212 index,
1213 block_type,
1214 id,
1215 name,
1216 type_name,
1217 data,
1218 origin,
1219 provider_data,
1220 } => self.start_block(
1221 index,
1222 block_type,
1223 id,
1224 name,
1225 type_name,
1226 data,
1227 origin,
1228 provider_data,
1229 )?,
1230 StreamEvent::TextDelta { index, text } => {
1231 match &mut self.block(index, BlockKind::text)?.kind {
1232 BlockKind::Text { text: buf, .. } => buf.push_str(&text),
1233 other => return Err(mismatch("text_delta", other, index)),
1234 }
1235 }
1236 StreamEvent::ReasoningDelta { index, text } => {
1237 match &mut self
1238 .block(index, || BlockKind::Reasoning(String::new()))?
1239 .kind
1240 {
1241 BlockKind::Reasoning(buf) => buf.push_str(&text),
1242 other => return Err(mismatch("reasoning_delta", other, index)),
1243 }
1244 }
1245 StreamEvent::RefusalDelta { index, text } => {
1246 match &mut self
1247 .block(index, || BlockKind::Refusal(String::new()))?
1248 .kind
1249 {
1250 BlockKind::Refusal(buf) => buf.push_str(&text),
1251 other => return Err(mismatch("refusal_delta", other, index)),
1252 }
1253 }
1254 StreamEvent::ToolCallDelta { index, arguments } => {
1255 let block = self.block(index, || BlockKind::ToolCall {
1256 id: None,
1257 name: None,
1258 arguments: String::new(),
1259 })?;
1260 match &mut block.kind {
1261 BlockKind::ToolCall { arguments: buf, .. }
1262 | BlockKind::ProviderToolCall { arguments: buf, .. } => {
1263 buf.push_str(&arguments)
1264 }
1265 other => return Err(mismatch("tool_call_delta", other, index)),
1266 }
1267 }
1268 StreamEvent::Citation { index, citation } => {
1269 match &mut self.block(index, BlockKind::text)?.kind {
1270 BlockKind::Text { citations, .. } => citations.push(citation),
1271 other => return Err(mismatch("citation", other, index)),
1272 }
1273 }
1274 StreamEvent::BlockProviderData {
1275 index,
1276 provider_data,
1277 } => {
1278 let block = self.open_block(index, "block_provider_data")?;
1279 block.provider_data.merge(provider_data);
1280 }
1281 StreamEvent::ResponseProviderData { provider_data } => {
1282 self.response.provider_data.merge(provider_data);
1283 }
1284 StreamEvent::ProviderDelta { .. } => {}
1285 StreamEvent::BlockStop { index } => return self.stop_block(index),
1286 StreamEvent::Warning { warning } => self.response.warnings.push(warning),
1287 StreamEvent::ResponseMetadata {
1288 finish_reason,
1289 finish_reason_raw,
1290 usage,
1291 usage_mode,
1292 id,
1293 model,
1294 request_id,
1295 provider,
1296 provider_data,
1297 } => {
1298 merge(&mut self.response.finish_reason, finish_reason);
1299 merge(&mut self.response.finish_reason_raw, finish_reason_raw);
1300 merge(&mut self.response.id, id);
1301 merge(&mut self.response.model, model);
1302 merge(&mut self.response.request_id, request_id);
1303 merge(&mut self.response.provider, provider);
1304 self.response.provider_data.merge(provider_data);
1305 if let Some(usage) = usage {
1306 match (usage_mode, &mut self.response.usage) {
1307 (UsageMetadataMode::Delta, Some(existing)) => *existing += usage,
1308 _ => self.response.usage = Some(usage),
1309 }
1310 }
1311 }
1312 StreamEvent::ResponseStop => self.saw_response_stop = true,
1313 }
1314 Ok(None)
1315 }
1316
1317 pub fn push_ref(&mut self, event: &StreamEvent) -> Result<Option<&ContentBlock>> {
1323 if self.metadata_only && event.is_block_event() && !self.saw_response_stop {
1325 return Ok(None);
1326 }
1327 self.push(event.clone())
1328 }
1329
1330 pub fn set_terminal_error(&mut self, error: Error) {
1332 self.terminal_error = Some(error);
1333 }
1334
1335 #[cfg_attr(not(feature = "extract"), allow(dead_code))]
1339 pub(crate) fn response_metadata(&self) -> &ChatResponse {
1340 &self.response
1341 }
1342
1343 #[must_use]
1352 pub fn block_type(&self, index: usize) -> Option<StreamBlockType> {
1353 match self.blocks.get(index)?.as_ref()? {
1354 Slot::Open(block) => Some(block.kind.block_type()),
1355 Slot::Closed { block_type, .. } => Some(*block_type),
1356 }
1357 }
1358
1359 pub fn finish(self) -> Result<ChatResponse> {
1374 let unidentified = self.has_unidentified_tool_call();
1375 let collected = self.build()?;
1376 if let Some(error) = collected.terminal_error {
1377 return Err(error);
1378 }
1379 let c = &collected.completeness;
1380 if !c.is_complete() {
1381 let hint = if unidentified {
1384 "; a tool call has no id or name: start it with \
1385 `StreamEvent::tool_call_start` (`provider_tool_call_start` for a \
1386 provider-executed one), not `block_start`"
1387 } else {
1388 ""
1389 };
1390 return Err(Error::stream(
1391 StreamErrorKind::Incomplete,
1392 format!(
1393 "stream incomplete: start={}, stop={}, open_blocks={}, dropped_incomplete_tool_calls={}{hint}",
1394 c.saw_response_start,
1395 c.saw_response_stop,
1396 c.has_open_blocks,
1397 c.dropped_incomplete_tool_calls
1398 ),
1399 ));
1400 }
1401 Ok(collected.response)
1402 }
1403
1404 pub fn finish_partial(self) -> Result<CollectedResponse> {
1410 self.build()
1411 }
1412
1413 fn block(
1416 &mut self,
1417 index: usize,
1418 default: impl FnOnce() -> BlockKind,
1419 ) -> Result<&mut BlockAccumulator> {
1420 self.reserve(index)?;
1421 let slot = self.blocks[index].get_or_insert_with(|| {
1422 Slot::Open(BlockAccumulator {
1423 kind: default(),
1424 provider_data: ProviderData::new(),
1425 started: false,
1426 })
1427 });
1428 match slot {
1429 Slot::Open(block) => Ok(block),
1430 Slot::Closed { block_type, .. } => Err(Error::stream_protocol(format!(
1431 "delta received for closed {} block at index {index}",
1432 block_type.label()
1433 ))),
1434 }
1435 }
1436
1437 fn reserve(&mut self, index: usize) -> Result<()> {
1443 if index >= self.blocks.len() {
1444 let gap = index - self.blocks.len();
1445 let skipped = self.skipped_indices.saturating_add(gap);
1446 if skipped > MAX_BLOCK_INDEX_GAP {
1447 return Err(Error::stream_protocol(format!(
1448 "block index {index} leaves {skipped} block indices unused in this \
1449 stream (at most {MAX_BLOCK_INDEX_GAP}); after {} blocks",
1450 self.blocks.len()
1451 )));
1452 }
1453 self.skipped_indices = skipped;
1454 self.blocks.resize_with(index + 1, || None);
1455 }
1456 Ok(())
1457 }
1458
1459 fn open_block(&mut self, index: usize, event: &str) -> Result<&mut BlockAccumulator> {
1461 match self.blocks.get_mut(index).and_then(Option::as_mut) {
1462 Some(Slot::Open(block)) => Ok(block),
1463 Some(Slot::Closed { block_type, .. }) => Err(Error::stream_protocol(format!(
1464 "{event} for closed {} block at index {index}",
1465 block_type.label()
1466 ))),
1467 None => Err(Error::stream_protocol(format!(
1468 "{event} for unknown index {index}"
1469 ))),
1470 }
1471 }
1472
1473 fn stop_block(&mut self, index: usize) -> Result<Option<&ContentBlock>> {
1474 let slot = self
1475 .blocks
1476 .get_mut(index)
1477 .and_then(Option::take)
1478 .ok_or_else(|| {
1479 Error::stream_protocol(format!("block_stop for unknown index {index}"))
1480 })?;
1481 let closed = match slot {
1482 Slot::Open(block) => Slot::Closed {
1483 block_type: block.kind.block_type(),
1484 block: block.finish(true),
1485 },
1486 Slot::Closed { block_type, block } => {
1487 self.blocks[index] = Some(Slot::Closed { block_type, block });
1488 return Err(Error::stream_protocol(format!(
1489 "duplicate block_stop for {} block at index {index}",
1490 block_type.label()
1491 )));
1492 }
1493 };
1494 let slot = self.blocks[index].insert(closed);
1495 Ok(match slot {
1496 Slot::Closed {
1497 block: Finished::Block(block),
1498 ..
1499 } => Some(block),
1500 _ => None,
1501 })
1502 }
1503
1504 #[allow(clippy::too_many_arguments)]
1505 fn start_block(
1506 &mut self,
1507 index: usize,
1508 block_type: StreamBlockType,
1509 id: Option<String>,
1510 name: Option<String>,
1511 type_name: Option<String>,
1512 data: Option<ExtraMap>,
1513 origin: Option<String>,
1514 provider_data: ProviderData,
1515 ) -> Result<()> {
1516 let fresh = match block_type {
1517 StreamBlockType::Text => BlockKind::text(),
1518 StreamBlockType::Reasoning => BlockKind::Reasoning(String::new()),
1519 StreamBlockType::Refusal => BlockKind::Refusal(String::new()),
1520 StreamBlockType::ToolCall => BlockKind::ToolCall {
1521 id: id.clone(),
1522 name: name.clone(),
1523 arguments: String::new(),
1524 },
1525 StreamBlockType::Image => {
1526 let (source, media_type) = image_from_data(data.clone())?;
1527 BlockKind::Image { source, media_type }
1528 }
1529 StreamBlockType::ProviderToolCall => BlockKind::ProviderToolCall {
1530 id: id.clone(),
1531 name: name.clone(),
1532 arguments: String::new(),
1533 origin,
1534 },
1535 StreamBlockType::ProviderToolResult => {
1536 let call_id = id.clone().ok_or_else(|| {
1537 Error::stream_protocol(format!(
1538 "provider_tool_result block at index {index} has no call id"
1539 ))
1540 })?;
1541 let mut data = data.clone().unwrap_or_default();
1542 BlockKind::ProviderToolResult {
1543 call_id,
1544 output: data.remove("output").unwrap_or_default(),
1545 is_error: data
1546 .remove("is_error")
1547 .and_then(|v| v.as_bool())
1548 .unwrap_or(false),
1549 origin,
1550 }
1551 }
1552 StreamBlockType::Other => BlockKind::Other {
1553 type_name: type_name.unwrap_or_else(|| "other".into()),
1554 data: data.unwrap_or_default(),
1555 origin,
1556 },
1557 };
1558
1559 self.reserve(index)?;
1560 let slot = &mut self.blocks[index];
1561 match slot {
1562 None => {
1563 *slot = Some(Slot::Open(BlockAccumulator {
1564 kind: fresh,
1565 provider_data,
1566 started: true,
1567 }));
1568 Ok(())
1569 }
1570 Some(Slot::Closed { block_type, .. }) => Err(Error::stream_protocol(format!(
1571 "block_start for closed {} block at index {index}",
1572 block_type.label()
1573 ))),
1574 Some(Slot::Open(block)) => {
1575 let label = block.kind.label();
1576 if block.started {
1577 return Err(Error::stream_protocol(format!(
1578 "duplicate block_start for {label} block at index {index}"
1579 )));
1580 }
1581 if label != fresh.label() {
1582 return Err(Error::stream_protocol(format!(
1583 "block_start for {} conflicted with existing {label} block at index {index}",
1584 fresh.label()
1585 )));
1586 }
1587 block.started = true;
1588 block.provider_data.merge(provider_data);
1589 if let BlockKind::ToolCall {
1590 id: existing_id,
1591 name: existing_name,
1592 ..
1593 } = &mut block.kind
1594 {
1595 merge_strict(existing_id, id, index, "id")?;
1596 merge_strict(existing_name, name, index, "name")?;
1597 }
1598 Ok(())
1599 }
1600 }
1601 }
1602
1603 fn has_unidentified_tool_call(&self) -> bool {
1605 self.blocks.iter().flatten().any(|slot| match slot {
1606 Slot::Open(block) => matches!(
1607 &block.kind,
1608 BlockKind::ToolCall { id, name, .. }
1609 | BlockKind::ProviderToolCall { id, name, .. }
1610 if id.is_none() || name.is_none()
1611 ),
1612 Slot::Closed { block, .. } => matches!(block, Finished::IncompleteToolCall),
1614 })
1615 }
1616
1617 fn build(self) -> Result<CollectedResponse> {
1618 let mut content = Vec::with_capacity(self.blocks.len());
1619 let mut has_open_blocks = false;
1620 let mut dropped_incomplete_tool_calls = false;
1621
1622 for slot in self.blocks.into_iter().flatten() {
1623 let finished = match slot {
1624 Slot::Open(block) => {
1625 has_open_blocks = true;
1626 block.finish(false)
1627 }
1628 Slot::Closed { block, .. } => block,
1629 };
1630 match finished {
1631 Finished::Block(block) => content.push(block),
1632 Finished::Empty => {}
1633 Finished::InvalidToolCall(call) => content.push(ContentBlock::ToolCall(call)),
1638 Finished::IncompleteToolCall => {
1639 dropped_incomplete_tool_calls = true;
1640 }
1641 }
1642 }
1643
1644 let mut response = self.response;
1645 response.content = content;
1646 Ok(CollectedResponse {
1647 response,
1648 completeness: StreamCompleteness {
1649 saw_response_start: self.saw_response_start,
1650 saw_response_stop: self.saw_response_stop,
1651 has_open_blocks,
1652 dropped_incomplete_tool_calls,
1653 },
1654 terminal_error: self.terminal_error,
1655 })
1656 }
1657}
1658
1659impl fmt::Debug for StreamCollector {
1661 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1662 f.debug_struct("StreamCollector")
1663 .field("blocks", &self.blocks.len())
1664 .field("saw_response_start", &self.saw_response_start)
1665 .field("saw_response_stop", &self.saw_response_stop)
1666 .field(
1667 "terminal_error",
1668 &self.terminal_error.as_ref().map(Error::kind_name),
1669 )
1670 .field("metadata_only", &self.metadata_only)
1671 .finish_non_exhaustive()
1672 }
1673}
1674
1675const MAX_BLOCK_INDEX_GAP: usize = 1024;
1679
1680fn merge<T>(slot: &mut Option<T>, value: Option<T>) {
1681 if value.is_some() {
1682 *slot = value;
1683 }
1684}
1685
1686fn merge_strict(
1687 existing: &mut Option<String>,
1688 incoming: Option<String>,
1689 index: usize,
1690 field: &str,
1691) -> Result<()> {
1692 match (existing.as_ref(), incoming) {
1693 (Some(current), Some(new)) if current != &new => Err(Error::stream_protocol(format!(
1694 "conflicting {field} for tool call block at index {index}"
1695 ))),
1696 (None, Some(new)) => {
1697 *existing = Some(new);
1698 Ok(())
1699 }
1700 _ => Ok(()),
1701 }
1702}
1703
1704fn mismatch(event: &str, kind: &BlockKind, index: usize) -> Error {
1705 Error::stream_protocol(format!(
1706 "{event} received for {} block at index {index}",
1707 kind.label()
1708 ))
1709}
1710
1711fn image_from_data(data: Option<ExtraMap>) -> Result<(MediaSource, Option<String>)> {
1714 let mut data = data.ok_or_else(|| Error::stream_protocol("image block missing source"))?;
1715 let source = data
1716 .remove("source")
1717 .ok_or_else(|| Error::stream_protocol("image block missing source"))?;
1718 let source = serde_json::from_value(source)
1719 .map_err(|e| Error::stream_decode(format!("image block source was invalid: {e}")))?;
1720 let media_type = match data.remove("media_type") {
1721 Some(serde_json::Value::String(media_type)) => Some(media_type),
1722 _ => None,
1723 };
1724 Ok((source, media_type))
1725}
1726
1727fn image_data(source: &MediaSource, media_type: Option<&str>) -> ExtraMap {
1728 let mut data = ExtraMap::new();
1729 data.insert(
1730 "source".into(),
1731 serde_json::to_value(source).expect("media source serializes"),
1732 );
1733 if let Some(media_type) = media_type {
1734 data.insert(
1735 "media_type".into(),
1736 serde_json::Value::String(media_type.to_owned()),
1737 );
1738 }
1739 data
1740}
1741
1742pub struct SingleResponseStream {
1746 events: std::vec::IntoIter<StreamEvent>,
1747}
1748
1749impl SingleResponseStream {
1750 #[must_use]
1752 pub fn new(response: ChatResponse) -> Self {
1753 Self {
1754 events: response
1755 .into_stream_events()
1756 .collect::<Vec<_>>()
1757 .into_iter(),
1758 }
1759 }
1760}
1761
1762impl fmt::Debug for SingleResponseStream {
1763 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1764 f.debug_struct("SingleResponseStream")
1765 .field("remaining_events", &self.events.len())
1766 .finish()
1767 }
1768}
1769
1770impl From<ChatResponse> for SingleResponseStream {
1771 fn from(response: ChatResponse) -> Self {
1772 Self::new(response)
1773 }
1774}
1775
1776impl Stream for SingleResponseStream {
1777 type Item = Result<StreamEvent>;
1778
1779 fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1780 Poll::Ready(self.events.next().map(Ok))
1781 }
1782}
1783
1784impl ChatResponse {
1785 #[must_use]
1787 pub fn into_stream(self) -> SingleResponseStream {
1788 SingleResponseStream::new(self)
1789 }
1790
1791 pub fn into_stream_events(self) -> impl Iterator<Item = StreamEvent> {
1793 let ChatResponse {
1794 content,
1795 finish_reason,
1796 finish_reason_raw,
1797 usage,
1798 model,
1799 id,
1800 request_id,
1801 provider,
1802 warnings,
1803 provider_data,
1804 finish_reported: _,
1805 } = self;
1806
1807 let mut events = Vec::with_capacity(content.len() * 3 + warnings.len() + 3);
1808 events.push(StreamEvent::ResponseStart {
1809 id: id.clone(),
1810 model: model.clone(),
1811 });
1812 events.extend(
1813 warnings
1814 .into_iter()
1815 .map(|warning| StreamEvent::Warning { warning }),
1816 );
1817 for (index, block) in content.into_iter().enumerate() {
1818 block_events(index, block, &mut events);
1819 }
1820 events.push(StreamEvent::ResponseMetadata {
1821 finish_reason,
1822 finish_reason_raw,
1823 usage,
1824 usage_mode: UsageMetadataMode::Snapshot,
1825 id,
1826 model,
1827 request_id,
1828 provider,
1829 provider_data,
1830 });
1831 events.push(StreamEvent::ResponseStop);
1832 events.into_iter()
1833 }
1834}
1835
1836fn block_events(index: usize, block: ContentBlock, events: &mut Vec<StreamEvent>) {
1837 let mut start = |block_type, id, name, type_name, data, origin, provider_data| {
1838 events.push(StreamEvent::BlockStart {
1839 index,
1840 block_type,
1841 id,
1842 name,
1843 type_name,
1844 data,
1845 origin,
1846 provider_data,
1847 });
1848 };
1849 match block {
1850 ContentBlock::Text(text) => {
1851 start(
1852 StreamBlockType::Text,
1853 None,
1854 None,
1855 None,
1856 None,
1857 None,
1858 text.provider_data,
1859 );
1860 if !text.text.is_empty() {
1861 events.push(StreamEvent::TextDelta {
1862 index,
1863 text: text.text,
1864 });
1865 }
1866 events.extend(
1867 text.citations
1868 .into_iter()
1869 .map(|citation| StreamEvent::Citation { index, citation }),
1870 );
1871 }
1872 ContentBlock::Image(image) => events.push(StreamEvent::image_start(index, image)),
1873 ContentBlock::Reasoning(reasoning) => {
1874 start(
1875 StreamBlockType::Reasoning,
1876 None,
1877 None,
1878 None,
1879 None,
1880 None,
1881 reasoning.provider_data,
1882 );
1883 if !reasoning.text.is_empty() {
1884 events.push(StreamEvent::ReasoningDelta {
1885 index,
1886 text: reasoning.text,
1887 });
1888 }
1889 }
1890 ContentBlock::Refusal(refusal) => {
1891 start(
1892 StreamBlockType::Refusal,
1893 None,
1894 None,
1895 None,
1896 None,
1897 None,
1898 refusal.provider_data,
1899 );
1900 if !refusal.text.is_empty() {
1901 events.push(StreamEvent::RefusalDelta {
1902 index,
1903 text: refusal.text,
1904 });
1905 }
1906 }
1907 ContentBlock::ToolCall(call) => {
1908 start(
1909 StreamBlockType::ToolCall,
1910 Some(call.id),
1911 Some(call.name),
1912 None,
1913 None,
1914 None,
1915 call.provider_data,
1916 );
1917 events.push(StreamEvent::ToolCallDelta {
1918 index,
1919 arguments: call.arguments,
1920 });
1921 }
1922 ContentBlock::ProviderToolCall(call) => {
1923 start(
1924 StreamBlockType::ProviderToolCall,
1925 Some(call.id),
1926 Some(call.name),
1927 None,
1928 None,
1929 call.origin,
1930 call.provider_data,
1931 );
1932 events.push(StreamEvent::ToolCallDelta {
1933 index,
1934 arguments: call.arguments,
1935 });
1936 }
1937 ContentBlock::ProviderToolResult(result) => {
1938 events.push(StreamEvent::provider_tool_result_start(index, result));
1939 }
1940 ContentBlock::Other(other) => events.push(StreamEvent::other_start(index, other)),
1941 }
1942 events.push(StreamEvent::BlockStop { index });
1943}
1944
1945pub trait ChatStreamExt: Stream<Item = Result<StreamEvent>> + Send {
1947 fn collect_response(self) -> impl std::future::Future<Output = Result<ChatResponse>> + Send;
1955
1956 fn next_event(
1979 &mut self,
1980 ) -> impl std::future::Future<Output = Option<Result<StreamEvent>>> + Send
1981 where
1982 Self: Unpin,
1983 {
1984 std::future::poll_fn(move |cx| Pin::new(&mut *self).poll_next(cx))
1985 }
1986
1987 fn with_collector(self) -> WithCollector<Self>
2025 where
2026 Self: Sized,
2027 {
2028 WithCollector {
2029 inner: Box::pin(self),
2030 collector: StreamCollector::new(),
2031 done: false,
2032 }
2033 }
2034}
2035
2036#[derive(Debug, Clone, PartialEq, Eq)]
2043#[non_exhaustive]
2044pub struct CollectedEvent {
2045 pub event: StreamEvent,
2047 pub completed: Option<ContentBlock>,
2061}
2062
2063impl CollectedEvent {
2064 #[must_use]
2067 pub fn new(event: StreamEvent, completed: Option<ContentBlock>) -> Self {
2068 Self { event, completed }
2069 }
2070
2071 #[must_use]
2073 pub fn into_parts(self) -> (StreamEvent, Option<ContentBlock>) {
2074 (self.event, self.completed)
2075 }
2076}
2077
2078impl From<CollectedEvent> for (StreamEvent, Option<ContentBlock>) {
2079 fn from(item: CollectedEvent) -> Self {
2080 item.into_parts()
2081 }
2082}
2083
2084pub struct WithCollector<S> {
2093 inner: Pin<Box<S>>,
2094 collector: StreamCollector,
2095 done: bool,
2096}
2097
2098impl<S> WithCollector<S>
2099where
2100 S: Stream<Item = Result<StreamEvent>>,
2101{
2102 pub async fn next_event(&mut self) -> Option<Result<CollectedEvent>> {
2107 std::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
2108 }
2109}
2110
2111impl<S> WithCollector<S> {
2112 #[must_use]
2114 pub fn collector(&self) -> &StreamCollector {
2115 &self.collector
2116 }
2117
2118 pub fn finish(self) -> Result<ChatResponse> {
2124 self.collector.finish()
2125 }
2126
2127 pub fn finish_partial(self) -> Result<CollectedResponse> {
2133 self.collector.finish_partial()
2134 }
2135}
2136
2137impl<S> fmt::Debug for WithCollector<S> {
2138 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2139 f.debug_struct("WithCollector")
2140 .field("collector", &self.collector)
2141 .field("done", &self.done)
2142 .finish_non_exhaustive()
2143 }
2144}
2145
2146impl<S> Stream for WithCollector<S>
2147where
2148 S: Stream<Item = Result<StreamEvent>>,
2149{
2150 type Item = Result<CollectedEvent>;
2151
2152 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2153 let this = &mut *self;
2154 if this.done {
2155 return Poll::Ready(None);
2156 }
2157 match this.inner.as_mut().poll_next(cx) {
2158 Poll::Ready(Some(Ok(event))) => match this.collector.push_ref(&event) {
2159 Ok(completed) => {
2160 let completed = completed.cloned();
2161 Poll::Ready(Some(Ok(CollectedEvent { event, completed })))
2162 }
2163 Err(error) => {
2164 this.done = true;
2165 this.collector.set_terminal_error(error.clone());
2166 Poll::Ready(Some(Err(error)))
2167 }
2168 },
2169 Poll::Ready(Some(Err(error))) => {
2170 this.done = true;
2171 this.collector.set_terminal_error(error.clone());
2172 Poll::Ready(Some(Err(error)))
2173 }
2174 Poll::Ready(None) => {
2175 this.done = true;
2176 Poll::Ready(None)
2177 }
2178 Poll::Pending => Poll::Pending,
2179 }
2180 }
2181}
2182
2183impl<S: Stream<Item = Result<StreamEvent>> + Send> ChatStreamExt for S {
2184 async fn collect_response(self) -> Result<ChatResponse> {
2185 use futures_util::StreamExt;
2186 let mut collector = StreamCollector::new();
2187 let mut stream = std::pin::pin!(self);
2188 while let Some(event) = stream.next().await {
2189 collector.push(event?)?;
2190 }
2191 collector.finish()
2192 }
2193}
2194
2195#[doc(alias = "CommittedStream")]
2206pub struct Prefixed<S> {
2207 buffer: VecDeque<StreamEvent>,
2208 inner: Option<S>,
2209}
2210
2211impl<S> Prefixed<S> {
2212 #[must_use]
2214 pub fn new(buffer: impl IntoIterator<Item = StreamEvent>, inner: Option<S>) -> Self {
2215 Self {
2216 buffer: buffer.into_iter().collect(),
2217 inner,
2218 }
2219 }
2220
2221 #[must_use]
2246 pub fn into_parts(self) -> (Vec<StreamEvent>, Option<S>) {
2247 (self.buffer.into(), self.inner)
2248 }
2249}
2250
2251impl<S> fmt::Debug for Prefixed<S> {
2252 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2253 f.debug_struct("Prefixed")
2254 .field("buffered_events", &self.buffer.len())
2255 .field("ended", &self.inner.is_none())
2256 .finish()
2257 }
2258}
2259
2260impl<S: Stream<Item = Result<StreamEvent>> + Unpin> Stream for Prefixed<S> {
2261 type Item = Result<StreamEvent>;
2262
2263 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2264 if let Some(event) = self.buffer.pop_front() {
2265 return Poll::Ready(Some(Ok(event)));
2266 }
2267 match self.inner.as_mut() {
2268 Some(inner) => Pin::new(inner).poll_next(cx),
2269 None => Poll::Ready(None),
2270 }
2271 }
2272}
2273
2274pub async fn until_committed<S>(mut stream: S) -> Result<Prefixed<S>>
2319where
2320 S: Stream<Item = Result<StreamEvent>> + Unpin,
2321{
2322 use futures_util::StreamExt;
2323 let mut buffered = Vec::new();
2324 loop {
2325 match stream.next().await {
2326 Some(Ok(event)) => {
2327 let committed = event.is_output();
2328 buffered.push(event);
2329 if committed {
2330 return Ok(Prefixed::new(buffered, Some(stream)));
2331 }
2332 }
2333 Some(Err(error)) => return Err(error),
2334 None => return Ok(Prefixed::new(buffered, None)),
2335 }
2336 }
2337}
2338
2339#[cfg(test)]
2340mod tests {
2341 use super::*;
2342 use crate::{CitationSource, RequestField, WarningKind};
2343 use futures_util::StreamExt;
2344 use serde_json::json;
2345
2346 #[tokio::test]
2349 async fn stream_types_debug_without_content() {
2350 let stream: ChatStream =
2351 Box::pin(ChatResponse::new(vec![ContentBlock::text("secret answer")]).into_stream());
2352 let single = SingleResponseStream::from(ChatResponse::new(vec![ContentBlock::text(
2353 "secret answer",
2354 )]));
2355 assert!(format!("{single:?}").starts_with("SingleResponseStream { remaining_events: "));
2356
2357 let mut collecting = stream.with_collector();
2358 while collecting.next_event().await.is_some() {}
2359 let debug = format!("{collecting:?}");
2360 assert!(
2361 debug.starts_with("WithCollector { collector: StreamCollector { blocks: 1"),
2362 "{debug}"
2363 );
2364 assert!(!debug.contains("secret"), "{debug}");
2365
2366 let prefixed = Prefixed::<ChatStream>::new([StreamEvent::text_delta(0, "secret")], None);
2367 assert_eq!(
2368 format!("{prefixed:?}"),
2369 "Prefixed { buffered_events: 1, ended: true }"
2370 );
2371 }
2372
2373 #[test]
2374 fn fieldless_enums_are_hashable() {
2375 use std::collections::HashSet;
2376
2377 use crate::{CapabilitySupport, ChatCapability, StreamErrorKind};
2378 let capabilities: HashSet<(ChatCapability, CapabilitySupport)> =
2379 [(ChatCapability::ToolCalls, CapabilitySupport::Supported)].into();
2380 assert!(capabilities.contains(&(ChatCapability::ToolCalls, CapabilitySupport::Supported)));
2381 let _: HashSet<StreamErrorKind> = HashSet::new();
2382 let _: HashSet<UsageMetadataMode> = HashSet::new();
2383 #[cfg(feature = "extract")]
2384 {
2385 let _: HashSet<crate::extract::RefusalSource> = HashSet::new();
2386 let _: HashSet<crate::extract::ExtractionMode> = HashSet::new();
2387 }
2388 }
2389
2390 fn sample_response() -> ChatResponse {
2391 let mut call = ToolCall::new("c1", "lookup", r#"{"q":1}"#);
2392 call.provider_data
2393 .insert_raw("google-ai", "thought_signature", json!("sig"));
2394 ChatResponse::new(vec![
2395 ContentBlock::Reasoning(Reasoning::new("think").with_provider_data(
2396 ProviderData::new().with_raw("anthropic", "signature", json!("s")),
2397 )),
2398 ContentBlock::Text(Text::new("hello").with_citations(vec![
2399 Citation::new(CitationSource::url("https://x", Some("X".into()))).with_span(0..5),
2400 ])),
2401 ContentBlock::Refusal(Refusal::new("I can't help with that.")),
2402 ContentBlock::ToolCall(call),
2403 ContentBlock::Image(Image::url("https://x/y.png")),
2404 ContentBlock::Other(
2405 OtherContent::new("citation", ExtraMap::new())
2406 .with_origin("anthropic")
2407 .with_provider_data(ProviderData::new().with_raw(
2408 "anthropic",
2409 "signature",
2410 json!("other-sig"),
2411 )),
2412 ),
2413 ])
2414 .finish_reason(FinishReason::ToolCalls)
2415 .finish_reason_raw("tool_use")
2416 .usage(Usage::new().input_tokens(1).output_tokens(2))
2417 .id("r1")
2418 .model("m")
2419 .request_id("req")
2420 .provider("demo")
2421 .warning(Warning::unsupported(RequestField::Seed, "ignored"))
2422 }
2423
2424 #[tokio::test]
2425 async fn response_round_trips_through_stream() {
2426 let response = sample_response();
2427 let collected = response
2428 .clone()
2429 .into_stream()
2430 .collect_response()
2431 .await
2432 .unwrap();
2433 assert_eq!(collected, response);
2434 }
2435
2436 #[tokio::test]
2439 async fn refusals_stream_apart_from_text() {
2440 let response = ChatResponse::new(vec![
2441 ContentBlock::text("Sure."),
2442 ContentBlock::Refusal(Refusal::new("I can't help with that.")),
2443 ]);
2444 let events: Vec<StreamEvent> = response.clone().into_stream_events().collect();
2445 assert!(events.iter().any(|event| matches!(
2446 event,
2447 StreamEvent::BlockStart {
2448 index: 1,
2449 block_type: StreamBlockType::Refusal,
2450 ..
2451 }
2452 )));
2453 assert!(events.contains(&StreamEvent::refusal_delta(1, "I can't help with that.")));
2454 assert!(!events.iter().any(|event| matches!(
2455 event,
2456 StreamEvent::TextDelta { text, .. } if text.contains("can't")
2457 )));
2458
2459 let collected = response
2460 .clone()
2461 .into_stream()
2462 .collect_response()
2463 .await
2464 .unwrap();
2465 assert_eq!(collected, response);
2466 assert_eq!(collected.first_refusal(), Some("I can't help with that."));
2467 assert_eq!(collected.text().as_deref(), Some("Sure."));
2468
2469 let mut collector = StreamCollector::new();
2471 collector
2472 .push(StreamEvent::refusal_delta(0, "No."))
2473 .unwrap();
2474 assert_eq!(collector.block_type(0), Some(StreamBlockType::Refusal));
2475 let error = collector.push(StreamEvent::text_delta(0, "x")).unwrap_err();
2476 assert!(error.to_string().contains("refusal block"), "{error}");
2477 }
2478
2479 #[test]
2480 fn all_block_types_listed() {
2481 fn _exhaustive(block_type: StreamBlockType) {
2484 match block_type {
2485 StreamBlockType::Text
2486 | StreamBlockType::Image
2487 | StreamBlockType::Reasoning
2488 | StreamBlockType::Refusal
2489 | StreamBlockType::ToolCall
2490 | StreamBlockType::ProviderToolCall
2491 | StreamBlockType::ProviderToolResult
2492 | StreamBlockType::Other => {}
2493 }
2494 }
2495 let unique: std::collections::HashSet<_> = StreamBlockType::ALL.iter().collect();
2496 assert_eq!(unique.len(), StreamBlockType::ALL.len());
2497 assert_eq!(StreamBlockType::ALL.len(), 8);
2498 }
2499
2500 #[test]
2501 fn refusals_are_output() {
2502 assert!(StreamEvent::refusal_delta(0, "No.").is_output());
2503 assert!(StreamEvent::block_start(0, StreamBlockType::Refusal).is_output());
2504 assert!(StreamEvent::refusal_delta(0, "No.").is_block_event());
2505 }
2506
2507 #[test]
2508 fn events_serialize_round_trip() {
2509 let events: Vec<StreamEvent> = sample_response().into_stream_events().collect();
2510 let value = serde_json::to_value(&events).unwrap();
2511 let back: Vec<StreamEvent> = serde_json::from_value(value).unwrap();
2512 assert_eq!(back, events);
2513 }
2514
2515 #[test]
2516 fn block_provider_data_merges() {
2517 let mut c = StreamCollector::new();
2518 c.push(StreamEvent::ResponseStart {
2519 id: None,
2520 model: None,
2521 })
2522 .unwrap();
2523 c.push(StreamEvent::block_start(0, StreamBlockType::Reasoning))
2524 .unwrap();
2525 c.push(StreamEvent::ReasoningDelta {
2526 index: 0,
2527 text: "a".into(),
2528 })
2529 .unwrap();
2530 c.push(StreamEvent::BlockProviderData {
2531 index: 0,
2532 provider_data: ProviderData::new().with_raw("anthropic", "signature", json!("s")),
2533 })
2534 .unwrap();
2535 c.push(StreamEvent::BlockStop { index: 0 }).unwrap();
2536 c.push(StreamEvent::Warning {
2537 warning: Warning::new(WarningKind::Dropped, crate::Subject::Response, "x"),
2538 })
2539 .unwrap();
2540 c.push(StreamEvent::ResponseStop).unwrap();
2541 let response = c.finish().unwrap();
2542 let reasoning = response.content[0].as_reasoning().unwrap();
2543 assert_eq!(
2544 reasoning.provider_data.get_raw("anthropic", "signature"),
2545 Some(&json!("s"))
2546 );
2547 assert_eq!(response.warnings.len(), 1);
2548 }
2549
2550 #[tokio::test]
2553 async fn other_block_keeps_its_provider_data() {
2554 let other = OtherContent::new("compaction", ExtraMap::new())
2555 .with_origin("anthropic")
2556 .with_provider_data(ProviderData::new().with_raw("anthropic", "start", json!(1)));
2557 let stream = futures_util::stream::iter(vec![
2558 Ok(StreamEvent::response_start(None, None)),
2559 Ok(StreamEvent::other_start(0, other)),
2560 Ok(StreamEvent::block_provider_data(
2561 0,
2562 ProviderData::new().with_raw("anthropic", "signature", json!("other-sig")),
2563 )),
2564 Ok(StreamEvent::block_stop(0)),
2565 Ok(StreamEvent::ResponseStop),
2566 ]);
2567 let response = stream.collect_response().await.unwrap();
2568 assert_eq!(
2569 response.content[0].provider_data(),
2570 &ProviderData::new()
2571 .with_raw("anthropic", "start", json!(1))
2572 .with_raw("anthropic", "signature", json!("other-sig"))
2573 );
2574 }
2575
2576 #[test]
2577 fn response_provider_data_merges_into_the_response() {
2578 let event = StreamEvent::response_provider_data(ProviderData::new().with_raw(
2579 "openai-compat:lumen",
2580 "search_results",
2581 json!([{"url": "https://a"}]),
2582 ));
2583 let value = serde_json::to_value(&event).unwrap();
2584 assert_eq!(
2585 value,
2586 json!({"event": "response_provider_data", "provider_data": {"openai-compat:lumen": {"search_results": [{"url": "https://a"}]}}})
2587 );
2588 assert_eq!(serde_json::from_value::<StreamEvent>(value).unwrap(), event);
2589 assert!(!event.is_output());
2590
2591 let mut c = StreamCollector::new();
2592 for item in sample_response().into_stream_events() {
2593 if matches!(item, StreamEvent::BlockStart { index: 0, .. }) {
2594 c.push(event.clone()).unwrap();
2595 }
2596 c.push(item).unwrap();
2597 }
2598 let response = c.finish().unwrap();
2599 assert_eq!(
2600 response
2601 .provider_data
2602 .get_raw("openai-compat:lumen", "search_results"),
2603 Some(&json!([{"url": "https://a"}]))
2604 );
2605 let mut expected = sample_response();
2606 expected.provider_data.insert_raw(
2607 "openai-compat:lumen",
2608 "search_results",
2609 json!([{"url": "https://a"}]),
2610 );
2611 assert_eq!(response, expected);
2612
2613 let mut c = StreamCollector::new();
2615 c.push(StreamEvent::ResponseStop).unwrap();
2616 let err = c
2617 .push(StreamEvent::response_provider_data(ProviderData::new()))
2618 .unwrap_err();
2619 assert!(err.to_string().contains("response_provider_data"), "{err}");
2620 }
2621
2622 #[test]
2623 fn provider_deltas_are_observed_not_collected() {
2624 let mut data = ExtraMap::new();
2625 data.insert("query".into(), json!("rust"));
2626 let delta = StreamEvent::provider_delta(
2627 3,
2628 "web_search_call.searching",
2629 data,
2630 Some("openai".into()),
2631 );
2632 let value = serde_json::to_value(&delta).unwrap();
2633 assert_eq!(
2634 value,
2635 json!({"event": "provider_delta", "index": 3, "type_name": "web_search_call.searching", "data": {"query": "rust"}, "origin": "openai"})
2636 );
2637 assert_eq!(serde_json::from_value::<StreamEvent>(value).unwrap(), delta);
2638 assert!(!delta.is_output());
2639
2640 let mut c = StreamCollector::new();
2641 for event in sample_response().into_stream_events() {
2642 if matches!(event, StreamEvent::ResponseStop) {
2643 c.push(delta.clone()).unwrap();
2644 }
2645 c.push(event).unwrap();
2646 }
2647 assert_eq!(c.finish().unwrap(), sample_response());
2648 }
2649
2650 #[test]
2651 fn push_returns_blocks_as_they_complete() {
2652 let mut c = StreamCollector::new();
2653 let mut completed = Vec::new();
2654 for event in sample_response().into_stream_events() {
2655 if let Some(block) = c.push(event).unwrap() {
2656 completed.push(block.clone());
2657 }
2658 }
2659 assert_eq!(completed, sample_response().content);
2660
2661 let mut c = StreamCollector::new();
2663 c.push(StreamEvent::block_start(0, StreamBlockType::Text))
2664 .unwrap();
2665 assert!(
2666 c.push(StreamEvent::BlockStop { index: 0 })
2667 .unwrap()
2668 .is_none()
2669 );
2670 c.push(StreamEvent::block_start(1, StreamBlockType::ToolCall))
2671 .unwrap();
2672 assert!(
2673 c.push(StreamEvent::BlockStop { index: 1 })
2674 .unwrap()
2675 .is_none()
2676 );
2677 c.push(StreamEvent::tool_call_start(2, "c", "f")).unwrap();
2678 c.push(StreamEvent::ToolCallDelta {
2679 index: 2,
2680 arguments: "{bad".into(),
2681 })
2682 .unwrap();
2683 assert!(
2684 c.push(StreamEvent::BlockStop { index: 2 })
2685 .unwrap()
2686 .is_none()
2687 );
2688 let partial = c.finish_partial().unwrap();
2691 let calls: Vec<_> = partial.response.tool_calls().collect();
2692 assert_eq!(calls, [&ToolCall::new("c", "f", "{bad")]);
2693 assert!(!calls[0].has_valid_arguments());
2694 assert!(partial.completeness.dropped_incomplete_tool_calls);
2695 }
2696
2697 #[test]
2698 fn a_tool_call_opened_without_an_id_names_its_constructor() {
2699 let mut c = StreamCollector::new();
2700 c.push(StreamEvent::response_start(Some("r".into()), None))
2701 .unwrap();
2702 c.push(StreamEvent::block_start(0, StreamBlockType::ToolCall))
2703 .unwrap();
2704 c.push(StreamEvent::BlockStop { index: 0 }).unwrap();
2705 c.push(StreamEvent::ResponseStop).unwrap();
2706 let err = c.finish().unwrap_err();
2707 assert!(
2708 err.message().contains("`StreamEvent::tool_call_start`"),
2709 "{err}"
2710 );
2711
2712 let mut c = StreamCollector::new();
2714 c.push(StreamEvent::response_start(Some("r".into()), None))
2715 .unwrap();
2716 c.push(StreamEvent::tool_call_start(0, "c", "f")).unwrap();
2717 let err = c.finish().unwrap_err();
2718 assert!(!err.message().contains("tool_call_start"), "{err}");
2719 }
2720
2721 #[test]
2722 fn unfinished_provider_tool_calls_are_dropped_from_partial_responses() {
2723 let mut c = StreamCollector::new();
2724 c.push(StreamEvent::provider_tool_call_start(
2725 0,
2726 "srv_1",
2727 "web_search",
2728 Some("anthropic".into()),
2729 ))
2730 .unwrap();
2731 c.push(StreamEvent::ToolCallDelta {
2732 index: 0,
2733 arguments: r#"{"query":"ru"#.into(),
2734 })
2735 .unwrap();
2736 let partial = c.finish_partial().unwrap();
2737 assert!(partial.response.content.is_empty());
2738 assert!(partial.completeness.dropped_incomplete_tool_calls);
2739 }
2740
2741 #[test]
2742 fn open_tool_calls_are_dropped_from_partial_responses() {
2743 let mut c = StreamCollector::new();
2744 c.push(StreamEvent::tool_call_start(0, "c1", "get_weather"))
2745 .unwrap();
2746 c.push(StreamEvent::tool_call_delta(0, r#"{"city":"Pa"#))
2747 .unwrap();
2748 let partial = c.finish_partial().unwrap();
2749 assert_eq!(partial.response.tool_calls().count(), 0);
2750 assert!(partial.completeness.dropped_incomplete_tool_calls);
2751 }
2752
2753 #[test]
2754 fn open_tool_call_without_id_or_name_has_a_block_type() {
2755 let mut c = StreamCollector::new();
2756 c.push(StreamEvent::tool_call_delta(0, "{")).unwrap();
2758 assert_eq!(c.block_type(0), Some(StreamBlockType::ToolCall));
2759 }
2760
2761 #[test]
2762 fn closed_tool_calls_complete_and_are_collected() {
2763 let mut c = StreamCollector::new();
2764 c.push(StreamEvent::tool_call_start(0, "c1", "get_weather"))
2765 .unwrap();
2766 c.push(StreamEvent::tool_call_delta(0, r#"{"city":"Paris"}"#))
2767 .unwrap();
2768 let first = c.push(StreamEvent::block_stop(0)).unwrap().cloned();
2769 c.push(StreamEvent::tool_call_start(1, "c2", "now"))
2770 .unwrap();
2771 let second = c.push(StreamEvent::block_stop(1)).unwrap().cloned();
2772 let expected = [
2773 ToolCall::new("c1", "get_weather", r#"{"city":"Paris"}"#),
2774 ToolCall::new("c2", "now", "{}"),
2775 ];
2776 assert_eq!(
2777 [first, second],
2778 expected
2779 .clone()
2780 .map(|call| Some(ContentBlock::ToolCall(call)))
2781 );
2782 let partial = c.finish_partial().unwrap();
2783 assert!(partial.response.tool_calls().eq(&expected));
2784 }
2785
2786 #[test]
2787 fn invalid_arguments_complete_as_nothing_but_are_collected() {
2788 let mut c = StreamCollector::new();
2789 c.push(StreamEvent::tool_call_start(0, "c1", "get_weather"))
2790 .unwrap();
2791 c.push(StreamEvent::tool_call_delta(0, r#"{"city":"#))
2792 .unwrap();
2793 assert!(c.push(StreamEvent::block_stop(0)).unwrap().is_none());
2794 let partial = c.finish_partial().unwrap();
2795 let call = partial.response.first_tool_call().unwrap();
2796 assert_eq!(
2797 (call.id.as_str(), call.arguments.as_str()),
2798 ("c1", r#"{"city":"#)
2799 );
2800 assert!(!call.has_valid_arguments());
2801 }
2802
2803 #[test]
2804 fn provider_tool_calls_never_appear_as_tool_calls() {
2805 let mut c = StreamCollector::new();
2806 c.push(StreamEvent::provider_tool_call_start(
2807 0,
2808 "srv_1",
2809 "web_search",
2810 Some("anthropic".into()),
2811 ))
2812 .unwrap();
2813 c.push(StreamEvent::tool_call_delta(0, r#"{"query":"rust"}"#))
2814 .unwrap();
2815 c.push(StreamEvent::block_stop(0)).unwrap();
2816
2817 c.push(StreamEvent::tool_call_start(1, "c1", "get_weather"))
2818 .unwrap();
2819 c.push(StreamEvent::block_stop(1)).unwrap();
2820 let partial = c.finish_partial().unwrap();
2821 let ids: Vec<_> = partial
2822 .response
2823 .tool_calls()
2824 .map(|call| call.id.as_str())
2825 .collect();
2826 assert_eq!(ids, ["c1"]);
2827 assert_eq!(partial.response.content.len(), 2);
2828 }
2829
2830 #[test]
2831 fn block_type_reports_open_and_closed_blocks() {
2832 let mut c = StreamCollector::new();
2833 assert_eq!(c.block_type(0), None);
2834 c.push(StreamEvent::text_delta(0, "Hi")).unwrap();
2835 c.push(StreamEvent::provider_tool_call_start(
2836 1,
2837 "srv_1",
2838 "web_search",
2839 None,
2840 ))
2841 .unwrap();
2842 c.push(StreamEvent::tool_call_start(3, "c1", "f")).unwrap();
2843 assert_eq!(c.block_type(0), Some(StreamBlockType::Text));
2844 assert_eq!(c.block_type(1), Some(StreamBlockType::ProviderToolCall));
2845 assert_eq!(c.block_type(2), None, "gap in the indexes");
2846 assert_eq!(c.block_type(3), Some(StreamBlockType::ToolCall));
2847 assert_eq!(c.block_type(4), None);
2848
2849 for index in [0, 1, 3] {
2850 c.push(StreamEvent::block_stop(index)).unwrap();
2851 }
2852 assert_eq!(c.block_type(0), Some(StreamBlockType::Text));
2853 assert_eq!(c.block_type(1), Some(StreamBlockType::ProviderToolCall));
2854 assert_eq!(c.block_type(2), None);
2855 assert_eq!(c.block_type(3), Some(StreamBlockType::ToolCall));
2856
2857 c.push(StreamEvent::block_start(5, StreamBlockType::Text))
2859 .unwrap();
2860 c.push(StreamEvent::block_stop(5)).unwrap();
2861 assert_eq!(c.block_type(5), Some(StreamBlockType::Text));
2862 }
2863
2864 #[test]
2865 fn stopped_blocks_reject_further_events() {
2866 let mut c = StreamCollector::new();
2867 c.push(StreamEvent::block_start(0, StreamBlockType::Text))
2868 .unwrap();
2869 c.push(StreamEvent::BlockStop { index: 0 }).unwrap();
2870 for event in [
2871 StreamEvent::TextDelta {
2872 index: 0,
2873 text: "late".into(),
2874 },
2875 StreamEvent::BlockProviderData {
2876 index: 0,
2877 provider_data: ProviderData::new().with_raw("x", "k", json!(1)),
2878 },
2879 StreamEvent::block_start(0, StreamBlockType::Text),
2880 StreamEvent::BlockStop { index: 0 },
2881 ] {
2882 let err = c.push(event).unwrap_err();
2883 assert!(err.to_string().contains("text block"), "{err}");
2884 }
2885 }
2886
2887 #[tokio::test]
2888 async fn with_collector_pairs_stops_with_blocks() {
2889 let mut stream = sample_response().into_stream().with_collector();
2890 let mut completed = Vec::new();
2891 while let Some(item) = stream.next_event().await {
2892 let (event, block) = item.unwrap().into_parts();
2893 assert_eq!(
2894 block.is_some(),
2895 matches!(event, StreamEvent::BlockStop { .. })
2896 );
2897 completed.extend(block);
2898 }
2899 assert_eq!(completed, sample_response().content);
2900 assert_eq!(stream.finish().unwrap(), sample_response());
2901 }
2902
2903 #[tokio::test]
2904 async fn with_collector_records_stream_errors() {
2905 let events = futures_util::stream::iter(vec![
2906 Ok(StreamEvent::ResponseStart {
2907 id: None,
2908 model: None,
2909 }),
2910 Ok(StreamEvent::TextDelta {
2911 index: 0,
2912 text: "partial".into(),
2913 }),
2914 Err(Error::overloaded("busy")),
2915 ]);
2916 let after = futures_util::stream::iter([Ok(StreamEvent::text_delta(0, "after"))]);
2917 let mut stream = events.chain(after).with_collector();
2918 let items: Vec<_> = stream.by_ref().collect().await;
2919 assert_eq!(items.len(), 3);
2921 assert!(items[2].is_err());
2922 let partial = stream.finish_partial().unwrap();
2923 assert_eq!(partial.response.text().as_deref(), Some("partial"));
2924 assert!(matches!(
2925 partial.terminal_error,
2926 Some(Error::Overloaded { .. })
2927 ));
2928 }
2929
2930 #[test]
2931 fn incomplete_stream_is_retryable() {
2932 let mut c = StreamCollector::new();
2933 c.push(StreamEvent::ResponseStart {
2934 id: None,
2935 model: None,
2936 })
2937 .unwrap();
2938 c.push(StreamEvent::TextDelta {
2939 index: 0,
2940 text: "partial".into(),
2941 })
2942 .unwrap();
2943 let err = c.finish().unwrap_err();
2944 assert!(matches!(
2945 err,
2946 Error::Stream {
2947 kind: StreamErrorKind::Incomplete,
2948 ..
2949 }
2950 ));
2951 assert!(err.is_retryable());
2952 }
2953
2954 #[test]
2955 fn protocol_violations_are_rejected() {
2956 let mut c = StreamCollector::new();
2957 c.push(StreamEvent::ResponseStop).unwrap();
2958 assert!(c.push(StreamEvent::ResponseStop).is_err());
2959
2960 let mut c = StreamCollector::new();
2961 c.push(StreamEvent::block_start(0, StreamBlockType::Text))
2962 .unwrap();
2963 assert!(
2964 c.push(StreamEvent::ToolCallDelta {
2965 index: 0,
2966 arguments: "{}".into()
2967 })
2968 .is_err()
2969 );
2970
2971 let mut c = StreamCollector::new();
2972 c.push(StreamEvent::tool_call_start(0, "a", "f")).unwrap();
2973 assert!(c.push(StreamEvent::tool_call_start(0, "b", "f")).is_err());
2974 }
2975
2976 #[test]
2977 fn invalid_tool_arguments_are_kept_like_unary() {
2978 let events = || {
2979 vec![
2980 StreamEvent::ResponseStart {
2981 id: None,
2982 model: None,
2983 },
2984 StreamEvent::tool_call_start(0, "c", "f"),
2985 StreamEvent::ToolCallDelta {
2986 index: 0,
2987 arguments: "{bad".into(),
2988 },
2989 StreamEvent::BlockStop { index: 0 },
2990 StreamEvent::ResponseStop,
2991 ]
2992 };
2993 let mut strict = StreamCollector::new();
2996 for e in events() {
2997 strict.push(e).unwrap();
2998 }
2999 let response = strict.finish().unwrap();
3000 let call = response.first_tool_call().unwrap();
3001 assert_eq!(call.arguments, "{bad");
3002 assert!(!call.has_valid_arguments());
3003
3004 let mut partial = StreamCollector::new();
3005 for e in events() {
3006 partial.push(e).unwrap();
3007 }
3008 let collected = partial.finish_partial().unwrap();
3009 assert!(!collected.completeness.dropped_incomplete_tool_calls);
3010 assert_eq!(collected.response.content.len(), 1);
3011 }
3012
3013 #[test]
3016 fn absurd_block_indices_are_protocol_errors() {
3017 for index in [usize::MAX, usize::MAX - 1, 10_000_000_000, 50_000_000] {
3018 let mut c = StreamCollector::new();
3019 let err = c.push(StreamEvent::text_delta(index, "x")).unwrap_err();
3020 assert!(
3021 matches!(
3022 err,
3023 Error::Stream {
3024 kind: StreamErrorKind::Protocol,
3025 ..
3026 }
3027 ),
3028 "{err:?}"
3029 );
3030 let mut c = StreamCollector::new();
3031 assert!(
3032 c.push(StreamEvent::block_start(index, StreamBlockType::Text))
3033 .is_err()
3034 );
3035 }
3036 let mut c = StreamCollector::new();
3038 c.push(StreamEvent::text_delta(MAX_BLOCK_INDEX_GAP, "x"))
3039 .unwrap();
3040 assert_eq!(
3041 c.block_type(MAX_BLOCK_INDEX_GAP),
3042 Some(StreamBlockType::Text)
3043 );
3044 }
3045
3046 #[test]
3050 fn stepped_block_indices_are_bounded_in_total() {
3051 let mut c = StreamCollector::new();
3052 let mut failed = None;
3053 for step in 1..=4 {
3054 let index = step * (MAX_BLOCK_INDEX_GAP / 2 + 1);
3055 if let Err(err) = c.push(StreamEvent::text_delta(index, "x")) {
3056 failed = Some((step, err));
3057 break;
3058 }
3059 }
3060 let (step, err) = failed.expect("stepped indices are refused");
3061 assert_eq!(step, 2, "the second gap passes the stream's total");
3062 assert!(
3063 matches!(
3064 err,
3065 Error::Stream {
3066 kind: StreamErrorKind::Protocol,
3067 ..
3068 }
3069 ),
3070 "{err:?}"
3071 );
3072 assert!(
3073 c.blocks.len() <= MAX_BLOCK_INDEX_GAP + 2,
3074 "{}",
3075 c.blocks.len()
3076 );
3077
3078 let mut c = StreamCollector::new();
3080 for index in 0..4 * MAX_BLOCK_INDEX_GAP {
3081 c.push(StreamEvent::text_delta(index, "x")).unwrap();
3082 }
3083 }
3084
3085 #[test]
3086 fn delta_usage_accumulates() {
3087 let mut c = StreamCollector::new();
3088 for usage in [Usage::new().output_tokens(2), Usage::new().output_tokens(3)] {
3089 c.push(StreamEvent::ResponseMetadata {
3090 usage: Some(usage),
3091 usage_mode: UsageMetadataMode::Delta,
3092 finish_reason: None,
3093 finish_reason_raw: None,
3094 id: None,
3095 model: None,
3096 request_id: None,
3097 provider: None,
3098 provider_data: ProviderData::new(),
3099 })
3100 .unwrap();
3101 }
3102 let collected = c.finish_partial().unwrap();
3103 assert_eq!(collected.response.usage.unwrap().output_tokens, Some(5));
3104 }
3105
3106 fn event_fixtures() -> Option<Vec<(String, Vec<StreamEvent>)>> {
3110 fn visit(dir: &std::path::Path, out: &mut Vec<std::path::PathBuf>) {
3111 let Ok(entries) = std::fs::read_dir(dir) else {
3112 return;
3113 };
3114 for path in entries.map(|entry| entry.unwrap().path()) {
3115 if path.is_dir() {
3116 visit(&path, out);
3117 } else if path.to_string_lossy().ends_with("_events.json") {
3118 out.push(path);
3119 }
3120 }
3121 }
3122 let crates = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
3123 let workspace = std::fs::read_to_string(crates.join("../Cargo.toml"))
3124 .is_ok_and(|manifest| manifest.contains("[workspace]"));
3125 if !workspace {
3126 eprintln!("skipped: the workspace stream fixtures are not available");
3127 return None;
3128 }
3129 let mut paths = Vec::new();
3130 for entry in std::fs::read_dir(&crates).unwrap() {
3131 visit(&entry.unwrap().path().join("fixtures"), &mut paths);
3132 }
3133 paths.sort();
3134 assert!(paths.len() >= 10, "found {} fixtures", paths.len());
3135 let fixtures = paths
3136 .into_iter()
3137 .map(|path| {
3138 let text = std::fs::read_to_string(&path).unwrap();
3139 let items: Vec<serde_json::Value> = serde_json::from_str(&text).unwrap();
3140 let events = items
3142 .into_iter()
3143 .take_while(|item| item.get("error").is_none())
3144 .map(|item| {
3145 serde_json::from_value(item)
3146 .unwrap_or_else(|err| panic!("{}: {err}", path.display()))
3147 })
3148 .collect();
3149 (path.display().to_string(), events)
3150 })
3151 .collect();
3152 Some(fixtures)
3153 }
3154
3155 fn collect_both(events: &[StreamEvent]) -> (CollectedResponse, CollectedResponse) {
3157 let mut full = StreamCollector::new();
3158 let mut metadata = StreamCollector::metadata_only();
3159 for event in events {
3160 let _ = full.push_ref(event);
3161 metadata.push_ref(event).unwrap();
3162 }
3163 (
3164 full.finish_partial().unwrap(),
3165 metadata.finish_partial().unwrap(),
3166 )
3167 }
3168
3169 fn assert_same_metadata(label: &str, events: &[StreamEvent]) {
3172 let (full, metadata) = collect_both(events);
3173 assert!(metadata.response.content.is_empty(), "{label}");
3174 assert_eq!(
3175 metadata.response.usage, full.response.usage,
3176 "{label}: usage"
3177 );
3178 assert_eq!(
3179 metadata.response.finish_reason, full.response.finish_reason,
3180 "{label}: finish reason"
3181 );
3182 assert_eq!(metadata.response.id, full.response.id, "{label}: id");
3183 assert_eq!(
3184 metadata.response.model, full.response.model,
3185 "{label}: model"
3186 );
3187 let mut expected = full.response;
3188 expected.content.clear();
3189 assert_eq!(metadata.response, expected, "{label}");
3190 }
3191
3192 #[test]
3193 fn metadata_only_collection_matches_full_collection_on_fixtures() {
3194 let Some(fixtures) = event_fixtures() else {
3195 return;
3196 };
3197 let mut with_usage = 0;
3198 for (label, events) in &fixtures {
3199 assert_same_metadata(label, events);
3200 if collect_both(events).1.response.usage.is_some() {
3201 with_usage += 1;
3202 }
3203 }
3204 assert!(with_usage > 0);
3205 }
3206
3207 #[test]
3208 fn metadata_only_collection_matches_full_collection_in_delta_usage_mode() {
3209 for (label, mut events) in event_fixtures().unwrap_or_default() {
3212 for event in &mut events {
3213 if let StreamEvent::ResponseMetadata { usage_mode, .. } = event {
3214 *usage_mode = UsageMetadataMode::Delta;
3215 }
3216 }
3217 assert_same_metadata(&format!("{label} (delta)"), &events);
3218 }
3219 }
3220
3221 #[test]
3222 fn metadata_only_collection_merges_metadata_and_skips_blocks() {
3223 let delta = |usage: Usage| {
3224 let mut event = StreamEvent::metadata();
3225 if let StreamEvent::ResponseMetadata {
3226 usage: slot,
3227 usage_mode,
3228 ..
3229 } = &mut event
3230 {
3231 *slot = Some(usage);
3232 *usage_mode = UsageMetadataMode::Delta;
3233 }
3234 event
3235 };
3236 let mut finish = StreamEvent::metadata();
3237 if let StreamEvent::ResponseMetadata {
3238 finish_reason,
3239 model,
3240 provider,
3241 ..
3242 } = &mut finish
3243 {
3244 *finish_reason = Some(FinishReason::Length);
3245 *model = Some("m-2".into());
3246 *provider = Some("demo".into());
3247 }
3248 let mut events = vec![
3249 StreamEvent::response_start(Some("r1".into()), Some("m".into())),
3250 delta(Usage::new().input_tokens(4).output_tokens(1)),
3251 ];
3252 events.extend(sample_response().into_stream_events().filter(|event| {
3253 !matches!(
3254 event,
3255 StreamEvent::ResponseStart { .. }
3256 | StreamEvent::ResponseMetadata { .. }
3257 | StreamEvent::ResponseStop
3258 )
3259 }));
3260 events.push(StreamEvent::text_delta(9, "dangling"));
3262 events.push(StreamEvent::block_stop(42));
3263 events.push(delta(Usage::new().output_tokens(2)));
3264 events.push(finish);
3265 events.push(StreamEvent::ResponseStop);
3266
3267 let mut collector = StreamCollector::metadata_only();
3268 for event in &events {
3269 assert!(collector.push_ref(event).unwrap().is_none());
3270 }
3271 assert!(
3272 collector
3273 .push_ref(&StreamEvent::text_delta(0, "late"))
3274 .is_err()
3275 );
3276 let collected = collector.finish_partial().unwrap();
3277 let response = &collected.response;
3278 assert!(response.content.is_empty());
3279 assert_eq!(
3280 response.usage,
3281 Some(Usage::new().input_tokens(4).output_tokens(3))
3282 );
3283 assert_eq!(response.finish_reason, Some(FinishReason::Length));
3284 assert_eq!(response.id.as_deref(), Some("r1"));
3285 assert_eq!(response.model.as_deref(), Some("m-2"));
3286 assert_eq!(response.provider.as_deref(), Some("demo"));
3287 assert_eq!(response.warnings, sample_response().warnings);
3288 assert!(collected.completeness.is_complete());
3289 }
3290
3291 #[tokio::test]
3292 async fn until_committed_buffers_and_replays() {
3293 let events = sample_response().into_stream();
3294 let prefixed = until_committed(events).await.unwrap();
3295 let all: Vec<_> = prefixed.map(|e| e.unwrap()).collect().await;
3296 assert_eq!(
3297 all,
3298 sample_response().into_stream_events().collect::<Vec<_>>()
3299 );
3300 }
3301
3302 #[tokio::test]
3303 async fn until_committed_reports_pre_output_errors() {
3304 let stream = futures_util::stream::iter(vec![
3305 Ok(StreamEvent::ResponseStart {
3306 id: None,
3307 model: None,
3308 }),
3309 Err(Error::overloaded("busy")),
3310 ]);
3311 let err = until_committed(stream).await.err().unwrap();
3312 assert!(matches!(err, Error::Overloaded { .. }));
3313 }
3314
3315 #[tokio::test]
3316 async fn until_committed_passes_post_output_errors_through() {
3317 let stream = futures_util::stream::iter(vec![
3318 Ok(StreamEvent::ResponseStart {
3319 id: None,
3320 model: None,
3321 }),
3322 Ok(StreamEvent::TextDelta {
3323 index: 0,
3324 text: "x".into(),
3325 }),
3326 Err(Error::overloaded("busy")),
3327 ]);
3328 let prefixed = until_committed(stream).await.unwrap();
3329 let items: Vec<_> = prefixed.collect().await;
3330 assert_eq!(items.len(), 3);
3331 assert!(items[2].is_err());
3332 }
3333
3334 #[tokio::test]
3335 async fn into_parts_returns_the_unyielded_prefix_and_the_rest() {
3336 let start = StreamEvent::response_start(None, None);
3337 let first = StreamEvent::text_delta(0, "a");
3338 let second = StreamEvent::text_delta(0, "b");
3339 let rest = futures_util::stream::iter(vec![Ok(second.clone())]);
3340 let mut prefixed = Prefixed::new([start.clone(), first.clone()], Some(rest));
3341
3342 assert_eq!(prefixed.next().await.unwrap().unwrap(), start);
3344 let (buffered, rest) = prefixed.into_parts();
3345 assert_eq!(buffered, std::slice::from_ref(&first));
3346
3347 let resumed: Vec<_> = Prefixed::new(buffered, rest)
3348 .map(Result::unwrap)
3349 .collect()
3350 .await;
3351 assert_eq!(resumed, [first, second]);
3352 }
3353
3354 #[tokio::test]
3355 async fn into_parts_of_an_ended_stream_has_no_rest() {
3356 let events = futures_util::stream::iter(vec![Ok(StreamEvent::response_start(None, None))]);
3357 let (buffered, rest) = until_committed(events).await.unwrap().into_parts();
3358 assert_eq!(buffered.len(), 1);
3359 assert!(rest.is_none());
3360 }
3361
3362 #[test]
3363 fn metadata_builder_sets_every_field_of_the_event() {
3364 let provider_data = ProviderData::new().with_raw("demo", "k", serde_json::json!(1));
3365 let usage = Usage::new().input_tokens(4).output_tokens(2);
3366 let built: StreamEvent = MetadataBuilder::new()
3367 .with_finish_reason(FinishReason::Length)
3368 .with_finish_reason_raw("max_tokens")
3369 .with_usage(usage.clone())
3370 .with_usage_mode(UsageMetadataMode::Delta)
3371 .with_id("r1".to_owned())
3372 .with_model("m")
3373 .with_request_id_opt(None)
3374 .with_provider("demo")
3375 .with_provider_data(provider_data.clone())
3376 .into();
3377
3378 let mut expected = StreamEvent::metadata();
3379 if let StreamEvent::ResponseMetadata {
3380 finish_reason,
3381 finish_reason_raw,
3382 usage: slot_usage,
3383 usage_mode,
3384 id,
3385 model,
3386 provider,
3387 provider_data: slot_provider_data,
3388 ..
3389 } = &mut expected
3390 {
3391 *finish_reason = Some(FinishReason::Length);
3392 *finish_reason_raw = Some("max_tokens".into());
3393 *slot_usage = Some(usage);
3394 *usage_mode = UsageMetadataMode::Delta;
3395 *id = Some("r1".into());
3396 *model = Some("m".into());
3397 *provider = Some("demo".into());
3398 *slot_provider_data = provider_data;
3399 }
3400 assert_eq!(built, expected);
3401 assert_eq!(MetadataBuilder::default().build(), StreamEvent::metadata());
3402 }
3403
3404 #[test]
3407 fn metadata_builder_opt_twins_keep_the_current_value_on_none() {
3408 let set = MetadataBuilder::new()
3409 .with_finish_reason(FinishReason::Stop)
3410 .with_finish_reason_raw("end_turn")
3411 .with_usage(Usage::new().output_tokens(1))
3412 .with_id("r1")
3413 .with_model("m")
3414 .with_request_id("req_1")
3415 .with_provider("demo");
3416 let kept = set
3417 .clone()
3418 .with_finish_reason_opt(None)
3419 .with_finish_reason_raw_opt(None)
3420 .with_usage_opt(None)
3421 .with_id_opt(None)
3422 .with_model_opt(None)
3423 .with_request_id_opt(None)
3424 .with_provider_opt(None);
3425 assert_eq!(kept, set);
3426
3427 let wire_id: Option<String> = Some("r2".to_owned());
3428 let StreamEvent::ResponseMetadata { id, .. } = set.with_id_opt(wire_id.as_deref()).build()
3429 else {
3430 unreachable!()
3431 };
3432 assert_eq!(id.as_deref(), Some("r2"));
3433
3434 let served_by = Some(Cow::Owned("compat:groq".to_owned()));
3435 let StreamEvent::ResponseMetadata { provider, .. } =
3436 MetadataBuilder::new().with_provider_opt(served_by).build()
3437 else {
3438 unreachable!()
3439 };
3440 assert_eq!(provider.as_deref(), Some("compat:groq"));
3441 }
3442
3443 #[test]
3446 fn image_media_type_survives_the_stream_events() {
3447 let image =
3448 Image::new(MediaSource::url("https://example.com/a")).with_media_type("image/jpeg");
3449 let response = ChatResponse::new(vec![ContentBlock::Image(image.clone())]);
3450 let mut collector = StreamCollector::new();
3451 for event in response.into_stream_events() {
3452 collector.push(event).unwrap();
3453 }
3454 let collected = collector.finish().unwrap();
3455 assert_eq!(collected.content, vec![ContentBlock::Image(image)]);
3456 }
3457}