1#[derive(Debug)]
14pub enum WireEvent<T> {
15 Known(T),
18 Unknown {
23 event_type: String,
25 value: crate::streaming::UnknownPayload,
28 },
29 Corrupt(serde_json::Error),
33}
34
35impl<T> WireEvent<T> {
36 pub fn map<U>(self, f: impl FnOnce(T) -> U) -> WireEvent<U> {
42 match self {
43 Self::Known(event) => WireEvent::Known(f(event)),
44 Self::Unknown { event_type, value } => WireEvent::Unknown { event_type, value },
45 Self::Corrupt(error) => WireEvent::Corrupt(error),
46 }
47 }
48}
49
50pub fn classify_tagged_frame<T>(
62 data: &str,
63 tag: &str,
64 is_known_event_type: impl Fn(&str) -> bool,
65) -> WireEvent<T>
66where
67 T: serde::de::DeserializeOwned,
68{
69 let scanned = match scan_discriminators(data, &[tag], true) {
70 Ok(scanned) => scanned,
71 Err(error) => return WireEvent::Corrupt(error),
72 };
73 match scanned {
74 DiscriminatorScan::Object(found) => {
75 match found.first().and_then(|key| key.string_value.as_deref()) {
76 Some(event_type) if !is_known_event_type(event_type) => {
77 unknown_with_value(data, event_type.to_owned())
78 }
79 _ => decode_known(data),
80 }
81 }
82 DiscriminatorScan::NotObject => unknown_with_value(data, String::new()),
87 }
88}
89
90fn unknown_with_value<T>(data: &str, event_type: String) -> WireEvent<T> {
93 match serde_json::from_str::<serde_json::Value>(data) {
94 Ok(value) => WireEvent::Unknown {
95 event_type,
96 value: value.into(),
97 },
98 Err(error) => WireEvent::Corrupt(error),
100 }
101}
102
103pub fn classify_chat_completions_frame<T>(data: &str) -> WireEvent<T>
112where
113 T: serde::de::DeserializeOwned,
114{
115 let scanned = match scan_discriminators(data, &["object", "choices"], true) {
116 Ok(scanned) => scanned,
117 Err(error) => return WireEvent::Corrupt(error),
118 };
119 let found = match scanned {
120 DiscriminatorScan::Object(found) => found,
121 DiscriminatorScan::NotObject => return unknown_with_value(data, String::new()),
125 };
126 let object_value = found.first().and_then(|key| key.string_value.as_deref());
127 let has_choices = found.get(1).is_some_and(|key| key.present);
128 let is_chat_chunk =
129 object_value.is_some_and(|object| object == "chat.completion.chunk") || has_choices;
130 if !is_chat_chunk {
131 return unknown_with_value(data, object_value.unwrap_or_default().to_owned());
132 }
133
134 decode_known(data)
135}
136
137pub fn classify_marker_keyed_frame<T>(data: &str, marker_keys: &[&str]) -> WireEvent<T>
146where
147 T: serde::de::DeserializeOwned,
148{
149 let scanned = match scan_discriminators(data, marker_keys, false) {
152 Ok(scanned) => scanned,
153 Err(error) => return WireEvent::Corrupt(error),
154 };
155 let recognizable = match &scanned {
156 DiscriminatorScan::Object(found) => found.iter().any(|key| key.present),
157 DiscriminatorScan::NotObject => false,
158 };
159 if !recognizable {
160 let value = match serde_json::from_str::<serde_json::Value>(data) {
164 Ok(value) => value,
165 Err(error) => return WireEvent::Corrupt(error),
167 };
168 let event_type = value
169 .as_object()
170 .map(|object| object.keys().cloned().collect::<Vec<_>>().join(","))
171 .unwrap_or_default();
172 return WireEvent::Unknown {
173 event_type,
174 value: value.into(),
175 };
176 }
177
178 decode_known(data)
179}
180
181pub fn classify_untyped_line<T>(line: &[u8]) -> WireEvent<T>
187where
188 T: serde::de::DeserializeOwned,
189{
190 match serde_json::from_slice::<T>(line) {
191 Ok(event) => WireEvent::Known(event),
192 Err(error) => WireEvent::Corrupt(error),
193 }
194}
195
196#[derive(Debug)]
200pub enum TypedEvent<T> {
201 Modeled(T),
203 Unrecognized {
206 event_type: String,
208 detail: String,
210 },
211 Malformed(String),
214}
215
216pub fn classify_typed_event<T>(event: TypedEvent<T>) -> WireEvent<T> {
226 match event {
227 TypedEvent::Modeled(event) => WireEvent::Known(event),
228 TypedEvent::Unrecognized { event_type, detail } => WireEvent::Unknown {
229 event_type,
230 value: serde_json::Value::String(detail).into(),
231 },
232 TypedEvent::Malformed(message) => {
233 WireEvent::Corrupt(<serde_json::Error as serde::de::Error>::custom(message))
234 }
235 }
236}
237
238pub fn classify_with_repair<T>(
250 data: &str,
251 classify: impl Fn(&str) -> WireEvent<T>,
252 repair: impl FnOnce(&str) -> Option<String>,
253 on_unrepairable: impl FnOnce(&serde_json::Error) -> serde_json::Error,
254 on_still_corrupt: impl FnOnce() -> serde_json::Error,
255) -> WireEvent<T> {
256 match classify(data) {
257 WireEvent::Corrupt(corrupt) => match repair(data) {
258 None => WireEvent::Corrupt(on_unrepairable(&corrupt)),
259 Some(repaired) => match classify(&repaired) {
260 WireEvent::Known(event) => WireEvent::Known(event),
261 WireEvent::Unknown { .. } | WireEvent::Corrupt(_) => {
265 WireEvent::Corrupt(on_still_corrupt())
266 }
267 },
268 },
269 event => event,
270 }
271}
272
273enum DiscriminatorScan {
291 Object(Vec<KeyScan>),
295 NotObject,
298}
299
300#[derive(Default, Clone)]
301struct KeyScan {
302 present: bool,
303 string_value: Option<String>,
304}
305
306fn scan_discriminators(
307 data: &str,
308 keys: &[&str],
309 reject_duplicates: bool,
310) -> Result<DiscriminatorScan, serde_json::Error> {
311 struct Scan<'a> {
312 keys: &'a [&'a str],
313 reject_duplicates: bool,
314 }
315
316 impl<'de> serde::de::Visitor<'de> for Scan<'_> {
317 type Value = DiscriminatorScan;
318
319 fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
320 formatter.write_str("a JSON value")
321 }
322
323 fn visit_map<A>(self, mut map: A) -> Result<DiscriminatorScan, A::Error>
324 where
325 A: serde::de::MapAccess<'de>,
326 {
327 let mut found = vec![KeyScan::default(); self.keys.len()];
328 while let Some(key) = map.next_key::<String>()? {
329 match self.keys.iter().position(|candidate| *candidate == key) {
330 Some(index) => {
331 let entry = found.get_mut(index).ok_or_else(|| {
332 serde::de::Error::custom("discriminator index out of range")
333 })?;
334 if entry.present {
335 if self.reject_duplicates {
336 return Err(serde::de::Error::custom(format!(
337 "duplicate `{key}` discriminator key in stream frame"
338 )));
339 }
340 map.next_value::<serde::de::IgnoredAny>()?;
343 continue;
344 }
345 entry.present = true;
346 entry.string_value = match map.next_value::<StringOrIgnored>()? {
349 StringOrIgnored::String(value) => Some(value),
350 StringOrIgnored::Ignored => None,
351 };
352 }
353 None => {
354 map.next_value::<serde::de::IgnoredAny>()?;
355 }
356 }
357 }
358 Ok(DiscriminatorScan::Object(found))
359 }
360
361 fn visit_bool<E>(self, _: bool) -> Result<DiscriminatorScan, E> {
363 Ok(DiscriminatorScan::NotObject)
364 }
365 fn visit_i64<E>(self, _: i64) -> Result<DiscriminatorScan, E> {
366 Ok(DiscriminatorScan::NotObject)
367 }
368 fn visit_u64<E>(self, _: u64) -> Result<DiscriminatorScan, E> {
369 Ok(DiscriminatorScan::NotObject)
370 }
371 fn visit_f64<E>(self, _: f64) -> Result<DiscriminatorScan, E> {
372 Ok(DiscriminatorScan::NotObject)
373 }
374 fn visit_str<E>(self, _: &str) -> Result<DiscriminatorScan, E> {
375 Ok(DiscriminatorScan::NotObject)
376 }
377 fn visit_unit<E>(self) -> Result<DiscriminatorScan, E> {
378 Ok(DiscriminatorScan::NotObject)
379 }
380 fn visit_seq<A>(self, mut seq: A) -> Result<DiscriminatorScan, A::Error>
381 where
382 A: serde::de::SeqAccess<'de>,
383 {
384 while seq.next_element::<serde::de::IgnoredAny>()?.is_some() {}
385 Ok(DiscriminatorScan::NotObject)
386 }
387 }
388
389 enum StringOrIgnored {
391 String(String),
392 Ignored,
393 }
394
395 impl<'de> serde::Deserialize<'de> for StringOrIgnored {
396 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
397 where
398 D: serde::Deserializer<'de>,
399 {
400 struct V;
401 impl<'de> serde::de::Visitor<'de> for V {
402 type Value = StringOrIgnored;
403 fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
404 formatter.write_str("any JSON value")
405 }
406 fn visit_str<E>(self, value: &str) -> Result<StringOrIgnored, E> {
407 Ok(StringOrIgnored::String(value.to_owned()))
408 }
409 fn visit_string<E>(self, value: String) -> Result<StringOrIgnored, E> {
410 Ok(StringOrIgnored::String(value))
411 }
412 fn visit_bool<E>(self, _: bool) -> Result<StringOrIgnored, E> {
413 Ok(StringOrIgnored::Ignored)
414 }
415 fn visit_i64<E>(self, _: i64) -> Result<StringOrIgnored, E> {
416 Ok(StringOrIgnored::Ignored)
417 }
418 fn visit_u64<E>(self, _: u64) -> Result<StringOrIgnored, E> {
419 Ok(StringOrIgnored::Ignored)
420 }
421 fn visit_f64<E>(self, _: f64) -> Result<StringOrIgnored, E> {
422 Ok(StringOrIgnored::Ignored)
423 }
424 fn visit_unit<E>(self) -> Result<StringOrIgnored, E> {
425 Ok(StringOrIgnored::Ignored)
426 }
427 fn visit_map<A>(self, mut map: A) -> Result<StringOrIgnored, A::Error>
428 where
429 A: serde::de::MapAccess<'de>,
430 {
431 while map
432 .next_entry::<serde::de::IgnoredAny, serde::de::IgnoredAny>()?
433 .is_some()
434 {}
435 Ok(StringOrIgnored::Ignored)
436 }
437 fn visit_seq<A>(self, mut seq: A) -> Result<StringOrIgnored, A::Error>
438 where
439 A: serde::de::SeqAccess<'de>,
440 {
441 while seq.next_element::<serde::de::IgnoredAny>()?.is_some() {}
442 Ok(StringOrIgnored::Ignored)
443 }
444 }
445 deserializer.deserialize_any(V)
446 }
447 }
448
449 let mut deserializer = serde_json::Deserializer::from_str(data);
450 let scanned = serde::Deserializer::deserialize_any(
451 &mut deserializer,
452 Scan {
453 keys,
454 reject_duplicates,
455 },
456 )?;
457 deserializer.end()?;
458 Ok(scanned)
459}
460
461fn decode_known<T>(data: &str) -> WireEvent<T>
462where
463 T: serde::de::DeserializeOwned,
464{
465 match serde_json::from_str::<T>(data) {
466 Ok(event) => WireEvent::Known(event),
467 Err(error) => WireEvent::Corrupt(error),
468 }
469}
470
471#[cfg(test)]
472mod tests {
473 use super::{WireEvent, classify_chat_completions_frame, classify_tagged_frame};
474
475 #[derive(Debug, serde::Deserialize)]
476 #[serde(tag = "type")]
477 enum TestEvent {
478 #[serde(rename = "text.delta")]
479 TextDelta { delta: String },
480 }
481
482 fn known(event_type: &str) -> bool {
483 event_type == "text.delta"
484 }
485
486 #[derive(Debug, serde::Deserialize)]
487 struct TestChunk {
488 #[allow(dead_code)]
489 choices: Vec<serde_json::Value>,
490 }
491
492 #[test]
493 fn tagged_known_frame_decodes() {
494 let event = classify_tagged_frame::<TestEvent>(
495 r#"{"type":"text.delta","delta":"hi"}"#,
496 "type",
497 known,
498 );
499 assert!(matches!(event, WireEvent::Known(TestEvent::TextDelta { delta }) if delta == "hi"));
500 }
501
502 #[test]
503 fn tagged_unknown_type_is_unknown() {
504 let event = classify_tagged_frame::<TestEvent>(r#"{"type":"future.event"}"#, "type", known);
505 assert!(matches!(
506 event,
507 WireEvent::Unknown { event_type, .. } if event_type == "future.event"
508 ));
509 }
510
511 #[test]
512 fn tagged_invalid_json_is_corrupt() {
513 let event = classify_tagged_frame::<TestEvent>("{not json", "type", known);
514 assert!(matches!(event, WireEvent::Corrupt(_)));
515 }
516
517 #[test]
518 fn tagged_known_type_with_defective_payload_is_corrupt() {
519 let event = classify_tagged_frame::<TestEvent>(
520 r#"{"type":"text.delta","delta":42}"#,
521 "type",
522 known,
523 );
524 assert!(matches!(event, WireEvent::Corrupt(_)));
525 }
526
527 #[test]
528 fn tagged_typeless_frame_is_corrupt() {
529 let event = classify_tagged_frame::<TestEvent>("{}", "type", known);
530 assert!(matches!(event, WireEvent::Corrupt(_)));
531 }
532
533 #[test]
538 fn tagged_duplicate_discriminator_is_corrupt() {
539 let event = classify_tagged_frame::<TestEvent>(
540 r#"{"type":"text.delta","type":"future.event","delta":"hi"}"#,
541 "type",
542 known,
543 );
544 assert!(matches!(event, WireEvent::Corrupt(_)));
545 }
546
547 #[test]
550 fn chat_duplicate_object_discriminator_is_corrupt() {
551 let event = classify_chat_completions_frame::<TestChunk>(
552 r#"{"object":"chat.completion.chunk","object":"future.thing","data":1}"#,
553 );
554 assert!(matches!(event, WireEvent::Corrupt(_)));
555 }
556
557 #[test]
558 fn chat_duplicate_choices_key_is_corrupt() {
559 let event = classify_chat_completions_frame::<TestChunk>(r#"{"choices":[],"choices":42}"#);
560 assert!(matches!(event, WireEvent::Corrupt(_)));
561 }
562
563 #[test]
564 fn tagged_duplicate_non_discriminator_key_still_classifies() {
565 let event = classify_tagged_frame::<TestEvent>(
569 r#"{"type":"text.delta","ignored":1,"ignored":2,"delta":"hi"}"#,
570 "type",
571 known,
572 );
573 assert!(matches!(event, WireEvent::Known(TestEvent::TextDelta { delta }) if delta == "hi"));
574 }
575
576 #[test]
577 fn chat_recognizable_chunk_decodes() {
578 let event = classify_chat_completions_frame::<TestChunk>(r#"{"choices":[]}"#);
579 assert!(matches!(event, WireEvent::Known(_)));
580 }
581
582 #[test]
583 fn chat_unrecognizable_json_is_unknown() {
584 let event = classify_chat_completions_frame::<TestChunk>(r#"{"object":"ping"}"#);
585 assert!(matches!(
586 event,
587 WireEvent::Unknown { event_type, .. } if event_type == "ping"
588 ));
589 }
590
591 #[test]
592 fn chat_recognizable_chunk_with_defective_payload_is_corrupt() {
593 let event = classify_chat_completions_frame::<TestChunk>(r#"{"choices":42}"#);
594 assert!(matches!(event, WireEvent::Corrupt(_)));
595 }
596
597 #[test]
598 fn chat_invalid_json_is_corrupt() {
599 let event = classify_chat_completions_frame::<TestChunk>("{not json");
600 assert!(matches!(event, WireEvent::Corrupt(_)));
601 }
602
603 #[test]
608 fn non_object_json_is_unknown_never_corrupt() {
609 for frame in ["null", "[]", "42", r#""ping""#] {
610 let event = classify_chat_completions_frame::<TestChunk>(frame);
611 assert!(
612 matches!(event, WireEvent::Unknown { .. }),
613 "chat classifier must skip {frame}, got {event:?}"
614 );
615 let event = classify_tagged_frame::<TestEvent>(frame, "type", known);
616 assert!(
617 matches!(event, WireEvent::Unknown { .. }),
618 "tagged classifier must skip {frame}, got {event:?}"
619 );
620 }
621 }
622
623 #[test]
624 fn tagged_dispatch_honors_a_non_type_tag_name() {
625 #[derive(Debug, serde::Deserialize)]
626 #[serde(tag = "event_type")]
627 enum EventTypeTagged {
628 #[serde(rename = "step.delta")]
629 StepDelta { delta: String },
630 }
631
632 let event = classify_tagged_frame::<EventTypeTagged>(
633 r#"{"event_type":"step.delta","delta":"hi"}"#,
634 "event_type",
635 |event_type| event_type == "step.delta",
636 );
637 assert!(matches!(
638 event,
639 WireEvent::Known(EventTypeTagged::StepDelta { delta }) if delta == "hi"
640 ));
641
642 let event = classify_tagged_frame::<EventTypeTagged>(
643 r#"{"event_type":"future.event"}"#,
644 "event_type",
645 |event_type| event_type == "step.delta",
646 );
647 assert!(matches!(
648 event,
649 WireEvent::Unknown { event_type, .. } if event_type == "future.event"
650 ));
651 }
652
653 #[test]
654 fn marker_keyed_recognizable_chunk_decodes() {
655 let event = super::classify_marker_keyed_frame::<TestChunk>(
656 r#"{"choices":[]}"#,
657 &["choices", "usage"],
658 );
659 assert!(matches!(event, WireEvent::Known(_)));
660 }
661
662 #[test]
663 fn marker_keyed_unrecognizable_json_is_unknown() {
664 let event = super::classify_marker_keyed_frame::<TestChunk>(
665 r#"{"noise":true,"other":1}"#,
666 &["choices", "usage"],
667 );
668 assert!(matches!(
669 event,
670 WireEvent::Unknown { event_type, .. } if event_type == "noise,other"
671 ));
672 }
673
674 #[test]
675 fn marker_keyed_recognizable_chunk_with_defective_payload_is_corrupt() {
676 let event =
677 super::classify_marker_keyed_frame::<TestChunk>(r#"{"choices":42}"#, &["choices"]);
678 assert!(matches!(event, WireEvent::Corrupt(_)));
679 }
680
681 #[test]
682 fn marker_keyed_invalid_json_is_corrupt() {
683 let event = super::classify_marker_keyed_frame::<TestChunk>("{not json", &["choices"]);
684 assert!(matches!(event, WireEvent::Corrupt(_)));
685 }
686
687 #[test]
688 fn typed_event_triage_maps_onto_the_shared_policy() {
689 let event = super::classify_typed_event(super::TypedEvent::Modeled(7u8));
691 assert!(matches!(event, WireEvent::Known(7)));
692
693 let event = super::classify_typed_event::<u8>(super::TypedEvent::Unrecognized {
696 event_type: "unknown".to_string(),
697 detail: "FutureEvent".to_string(),
698 });
699 assert!(matches!(
700 event,
701 WireEvent::Unknown { event_type, value }
702 if event_type == "unknown" && value.value() == &serde_json::Value::String("FutureEvent".into())
703 ));
704
705 let event =
707 super::classify_typed_event::<u8>(super::TypedEvent::Malformed("bad frame".into()));
708 assert!(
709 matches!(event, WireEvent::Corrupt(error) if error.to_string().contains("bad frame"))
710 );
711 }
712
713 #[test]
714 fn untyped_line_is_known_or_corrupt() {
715 assert!(matches!(
716 super::classify_untyped_line::<TestChunk>(br#"{"choices":[]}"#),
717 WireEvent::Known(_)
718 ));
719 assert!(matches!(
720 super::classify_untyped_line::<TestChunk>(b"{not json"),
721 WireEvent::Corrupt(_)
722 ));
723 }
724}