Skip to main content

agent_framework_core/
streaming.rs

1//! Utilities for decoding streamed bytes.
2
3/// An incremental UTF-8 decoder for byte streams.
4///
5/// Network streams (SSE and friends) deliver bytes in arbitrary chunks, so a
6/// multi-byte UTF-8 character can be split across two chunks. Decoding each
7/// chunk independently with `String::from_utf8_lossy` corrupts such
8/// characters: the truncated head becomes U+FFFD and the continuation bytes
9/// in the next chunk become more U+FFFD. This decoder holds back an
10/// incomplete trailing sequence until the rest of it arrives.
11///
12/// Genuinely invalid bytes (not a truncated tail) are replaced with U+FFFD,
13/// matching lossy semantics.
14#[derive(Debug, Default)]
15pub struct Utf8StreamDecoder {
16    pending: Vec<u8>,
17}
18
19impl Utf8StreamDecoder {
20    /// Create an empty decoder.
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    /// Feed `bytes`, returning all completely-decodable text.
26    pub fn push(&mut self, bytes: &[u8]) -> String {
27        self.pending.extend_from_slice(bytes);
28        let mut out = String::new();
29        loop {
30            match std::str::from_utf8(&self.pending) {
31                Ok(s) => {
32                    out.push_str(s);
33                    self.pending.clear();
34                    break;
35                }
36                Err(e) => {
37                    let valid = e.valid_up_to();
38                    // Safety of unwrap: `valid_up_to` guarantees validity.
39                    out.push_str(std::str::from_utf8(&self.pending[..valid]).unwrap());
40                    match e.error_len() {
41                        // Definitely-invalid sequence: replace and continue.
42                        Some(n) => {
43                            out.push('\u{FFFD}');
44                            self.pending.drain(..valid + n);
45                        }
46                        // Incomplete trailing sequence: keep it for the next
47                        // chunk.
48                        None => {
49                            self.pending.drain(..valid);
50                            break;
51                        }
52                    }
53                }
54            }
55        }
56        out
57    }
58
59    /// Drain any held-back bytes at end of stream (lossy).
60    pub fn flush(&mut self) -> String {
61        let out = String::from_utf8_lossy(&self.pending).into_owned();
62        self.pending.clear();
63        out
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn passes_ascii_through() {
73        let mut d = Utf8StreamDecoder::new();
74        assert_eq!(d.push(b"hello"), "hello");
75        assert_eq!(d.flush(), "");
76    }
77
78    #[test]
79    fn reassembles_multibyte_char_split_across_chunks() {
80        let mut d = Utf8StreamDecoder::new();
81        let euro = "€".as_bytes(); // 3 bytes: E2 82 AC
82        assert_eq!(d.push(&euro[..1]), "");
83        assert_eq!(d.push(&euro[1..2]), "");
84        assert_eq!(d.push(&euro[2..]), "€");
85    }
86
87    #[test]
88    fn mixed_text_and_split_emoji() {
89        let mut d = Utf8StreamDecoder::new();
90        let s = "ok 🚀 done".as_bytes();
91        let cut = 5; // inside the 4-byte emoji (starts at index 3)
92        let first = d.push(&s[..cut]);
93        let second = d.push(&s[cut..]);
94        assert_eq!(format!("{first}{second}"), "ok 🚀 done");
95    }
96
97    #[test]
98    fn replaces_genuinely_invalid_bytes() {
99        let mut d = Utf8StreamDecoder::new();
100        assert_eq!(d.push(&[b'a', 0xFF, b'b']), "a\u{FFFD}b");
101    }
102
103    #[test]
104    fn flush_lossily_drains_incomplete_tail() {
105        let mut d = Utf8StreamDecoder::new();
106        assert_eq!(d.push("€".as_bytes()[..2].as_ref()), "");
107        assert_eq!(d.flush(), "\u{FFFD}");
108    }
109}