procyon 0.1.2

//! Server-sent event framing and reading, shared by both provider adapters.
//!
//! Both dialects deliver a stream as SSE, and neither guarantees that a TCP chunk boundary falls
//! on a line boundary: one `data:` line can arrive split across two reads, and one read can carry
//! several events. Keeping the reassembly here means it is tested once against those splits
//! rather than reimplemented per adapter.
//!
//! [`EventReader`] carries the other half both adapters had written out longhand: read a chunk,
//! frame it, hand each payload to the state machine, drain the decoder when the connection ends,
//! and stop early when the state machine says to. The *event grammars* genuinely differ — one is
//! indexed block deltas, the other is choice deltas — so what is shared is the reading, and each
//! adapter keeps its own [`EventSink`].

use color_eyre::Result;
use futures_util::StreamExt as _;

/// Accumulates stream bytes and yields the payload of each complete `data:` line.
#[derive(Debug, Default)]
pub struct SseDecoder {
    buffer: String,
    /// Bytes of a UTF-8 sequence split across chunk boundaries, awaiting its continuation.
    pending: Vec<u8>,
}

impl SseDecoder {
    pub fn new() -> Self {
        Self::default()
    }

    /// Feeds one chunk of raw stream bytes and returns the `data:` payloads it completed.
    ///
    /// A chunk boundary can fall inside a multi-byte character, so the incomplete tail is held
    /// for the next chunk instead of being replaced: decoding each chunk on its own turned every
    /// accented character unlucky enough to straddle a read into `U+FFFD`. Bytes that are
    /// genuinely not UTF-8 are still replaced, but recorded rather than dropped silently — a
    /// corrupted response otherwise reads as a model that wrote nonsense.
    pub fn push_bytes(&mut self, chunk: &[u8]) -> Vec<String> {
        self.pending.extend_from_slice(chunk);
        let text = self.take_decodable();
        self.push(&text)
    }

    /// Drains every byte of `pending` that forms complete characters, leaving an incomplete
    /// trailing sequence in place.
    fn take_decodable(&mut self) -> String {
        let mut out = String::new();

        loop {
            match std::str::from_utf8(&self.pending) {
                Ok(valid) => {
                    out.push_str(valid);
                    self.pending.clear();
                    return out;
                }
                Err(e) => {
                    let good = e.valid_up_to();
                    // Everything before the offending byte is valid by definition.
                    out.push_str(std::str::from_utf8(&self.pending[..good]).unwrap_or_default());

                    match e.error_len() {
                        // Truncated, not invalid: the rest may arrive in the next chunk.
                        None => {
                            self.pending.drain(..good);
                            return out;
                        }
                        Some(len) => {
                            crate::diag::warn(format!(
                                "stream: replaced {} byte(s) of invalid UTF-8 at offset {}",
                                len, good
                            ));
                            out.push('\u{FFFD}');
                            self.pending.drain(..good + len);
                        }
                    }
                }
            }
        }
    }

    /// Feeds one chunk and returns the `data:` payloads it completed, in order.
    ///
    /// A partial trailing line is retained for the next call. Comment lines (`:`), blank lines
    /// and non-`data` fields (`event:`, `id:`, `retry:`) are dropped: no provider here carries
    /// meaning in them, and an unknown field must not be mistaken for a payload.
    pub fn push(&mut self, chunk: &str) -> Vec<String> {
        self.buffer.push_str(chunk);
        let mut payloads = Vec::new();

        while let Some(newline) = self.buffer.find('\n') {
            let line = self.buffer[..newline].to_string();
            self.buffer.drain(..newline + 1);
            if let Some(payload) = Self::payload_of(&line) {
                payloads.push(payload);
            }
        }

        payloads
    }

    /// Yields a final payload left unterminated by the end of the stream.
    ///
    /// A well-formed stream ends with a newline and this returns `None`. A provider or proxy that
    /// closes the connection right after the last event would otherwise have that event silently
    /// dropped, which reads as an empty response rather than as the truncation it is.
    pub fn flush(&mut self) -> Option<String> {
        // A sequence still pending when the stream ends will never be completed. It is one
        // character of a payload the caller is about to parse, so it is marked, not hidden.
        if !self.pending.is_empty() {
            crate::diag::warn(format!(
                "stream: ended mid-character, {} byte(s) unterminated",
                self.pending.len()
            ));
            self.pending.clear();
            self.buffer.push('\u{FFFD}');
        }

        let line = std::mem::take(&mut self.buffer);
        Self::payload_of(&line)
    }

    fn payload_of(line: &str) -> Option<String> {
        // `trim` also removes the `\r` of CRLF framing, which some proxies introduce.
        let line = line.trim();

        if line.is_empty() || line.starts_with(':') {
            return None;
        }

        // The space after the colon is optional in the SSE grammar, and local servers do emit the
        // bare form. Requiring it would drop every event from those.
        let payload = line
            .strip_prefix("data:")
            .map(|rest| rest.strip_prefix(' ').unwrap_or(rest))?;

        Some(payload.to_string())
    }
}

/// Whether the reader should keep pulling from the stream.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Flow {
    Continue,
    /// The sink has everything it needs — typically because the consumer hung up.
    Stop,
}

/// A provider's stream state machine: one payload in, a decision out.
///
/// Absorbing is a plain function of state, which is what lets a whole stream be driven in a test
/// without a socket.
pub trait EventSink {
    fn absorb(&mut self, payload: &str) -> Result<Flow>;
}

/// Frames a byte stream into events and feeds them to a sink.
pub struct EventReader<S> {
    decoder: SseDecoder,
    sink: S,
}

impl<S: EventSink> EventReader<S> {
    pub fn new(sink: S) -> Self {
        Self {
            decoder: SseDecoder::new(),
            sink,
        }
    }

    /// Reads the response to its end, or until the sink asks to stop.
    pub async fn read(&mut self, response: reqwest::Response) -> Result<()> {
        let mut stream = response.bytes_stream();

        while let Some(chunk) = stream.next().await {
            if self.feed(&chunk?)? == Flow::Stop {
                return Ok(());
            }
        }

        // The connection closing is the only signal that a trailing event will never be
        // terminated, so the decoder is drained here rather than left holding it.
        self.finish()
    }

    /// Feeds one chunk of raw bytes, returning whether to keep reading.
    pub fn feed(&mut self, chunk: &[u8]) -> Result<Flow> {
        for payload in self.decoder.push_bytes(chunk) {
            if self.sink.absorb(&payload)? == Flow::Stop {
                return Ok(Flow::Stop);
            }
        }
        Ok(Flow::Continue)
    }

    /// Hands the sink any event the end of the stream left unterminated.
    pub fn finish(&mut self) -> Result<()> {
        if let Some(payload) = self.decoder.flush() {
            self.sink.absorb(&payload)?;
        }
        Ok(())
    }

    pub fn into_sink(self) -> S {
        self.sink
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn all(chunks: &[&str]) -> Vec<String> {
        let mut decoder = SseDecoder::new();
        let mut out: Vec<String> = chunks.iter().flat_map(|c| decoder.push(c)).collect();
        out.extend(decoder.flush());
        out
    }

    #[test]
    fn decodes_one_event_per_line() {
        assert_eq!(all(&["data: a\ndata: b\n"]), vec!["a", "b"]);
    }

    // The likeliest framing bug: a chunk boundary inside a line. Splitting the same stream at
    // every possible offset must always yield the same events.
    #[test]
    fn any_chunk_boundary_yields_the_same_events() {
        let stream = "data: {\"i\":1}\n\ndata: {\"i\":2}\n\ndata: [DONE]\n\n";
        let expected = vec!["{\"i\":1}", "{\"i\":2}", "[DONE]"];

        for split in 0..stream.len() {
            let (head, tail) = stream.split_at(split);
            assert_eq!(all(&[head, tail]), expected, "split at {}", split);
        }
    }

    #[test]
    fn one_chunk_may_carry_many_events() {
        assert_eq!(
            all(&["data: a\n\ndata: b\n\ndata: c\n\n"]),
            vec!["a", "b", "c"]
        );
    }

    #[test]
    fn a_payload_may_arrive_one_byte_at_a_time() {
        let chunks: Vec<String> = "data: hello\n".chars().map(|c| c.to_string()).collect();
        let refs: Vec<&str> = chunks.iter().map(String::as_str).collect();

        assert_eq!(all(&refs), vec!["hello"]);
    }

    // Requiring the space would drop every event from a server that emits the bare form.
    #[test]
    fn the_space_after_the_colon_is_optional() {
        assert_eq!(all(&["data:a\ndata: b\n"]), vec!["a", "b"]);
    }

    #[test]
    fn crlf_framing_decodes() {
        assert_eq!(all(&["data: a\r\ndata: b\r\n"]), vec!["a", "b"]);
    }

    #[test]
    fn comments_blank_lines_and_other_fields_are_dropped() {
        assert_eq!(
            all(&[": keepalive\n\nevent: message\nid: 7\nretry: 100\ndata: a\n"]),
            vec!["a"]
        );
    }

    // Without the flush the last event of a truncated stream disappears, which the caller reads
    // as a short-but-successful response instead of a cut-off one.
    #[test]
    fn a_final_line_without_a_newline_is_still_yielded() {
        assert_eq!(all(&["data: a\ndata: b"]), vec!["a", "b"]);
    }

    #[test]
    fn a_well_formed_stream_flushes_nothing() {
        let mut decoder = SseDecoder::new();
        assert_eq!(decoder.push("data: a\n"), vec!["a"]);
        assert_eq!(decoder.flush(), None);
    }

    // JSON payloads contain colons and braces; only the line prefix may be treated as framing.
    #[test]
    fn json_payloads_pass_through_untouched() {
        let payload = r#"{"a": {"b": ":data: not-a-line"}}"#;
        assert_eq!(all(&[&format!("data: {}\n", payload)]), vec![payload]);
    }

    #[test]
    fn multibyte_content_survives_a_split_between_its_bytes() {
        let stream = "data: {\"t\":\"ação — ok\"}\n";
        // Splitting inside a multibyte sequence is only representable at the byte level, which is
        // where the adapter's lossy UTF-8 conversion already happens; this pins the char-boundary
        // case the decoder itself must handle.
        for split in 0..stream.chars().count() {
            let head: String = stream.chars().take(split).collect();
            let tail: String = stream.chars().skip(split).collect();
            assert_eq!(all(&[&head, &tail]), vec!["{\"t\":\"ação — ok\"}"]);
        }
    }

    fn all_bytes(chunks: &[&[u8]]) -> Vec<String> {
        let mut decoder = SseDecoder::new();
        let mut out: Vec<String> = chunks.iter().flat_map(|c| decoder.push_bytes(c)).collect();
        out.extend(decoder.flush());
        out
    }

    // The byte-level counterpart of the case above, and the one the adapters actually hit: a read
    // may end anywhere, including between the two bytes of `ç`.
    #[test]
    fn any_byte_boundary_yields_the_same_payload() {
        let stream = "data: {\"t\":\"ação — ok\"}\n".as_bytes();
        let expected = vec!["{\"t\":\"ação — ok\"}"];

        for split in 0..stream.len() {
            let (head, tail) = stream.split_at(split);
            assert_eq!(all_bytes(&[head, tail]), expected, "split at {}", split);
        }
    }

    #[test]
    fn a_multibyte_character_may_arrive_one_byte_at_a_time() {
        let chunks: Vec<Vec<u8>> = "data: — ção\n".bytes().map(|b| vec![b]).collect();
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        assert_eq!(all_bytes(&refs), vec!["— ção"]);
    }

    // Invalid bytes still become U+FFFD — the payload must stay parseable — but the substitution
    // is recorded, which is the difference between a diagnosable corruption and a mystery.
    #[test]
    fn invalid_bytes_are_replaced_and_recorded() {
        let _guard = crate::diag::test_lock();
        crate::diag::drain();

        let mut stream = b"data: a".to_vec();
        stream.push(0xFF);
        stream.extend_from_slice(b"b\n");

        assert_eq!(all_bytes(&[&stream]), vec!["a\u{FFFD}b"]);
        let warnings = crate::diag::drain();
        assert!(
            warnings.iter().any(|w| w.contains("invalid UTF-8")),
            "expected a recorded warning, got {:?}",
            warnings
        );
    }

    /// Records what it was given, and stops on a payload of `halt`.
    #[derive(Default)]
    struct Recorder {
        seen: Vec<String>,
    }

    impl EventSink for Recorder {
        fn absorb(&mut self, payload: &str) -> Result<Flow> {
            self.seen.push(payload.to_string());
            Ok(if payload == "halt" {
                Flow::Stop
            } else {
                Flow::Continue
            })
        }
    }

    fn read_all(chunks: &[&[u8]]) -> Vec<String> {
        let mut reader = EventReader::new(Recorder::default());
        for chunk in chunks {
            if reader.feed(chunk).unwrap() == Flow::Stop {
                return reader.into_sink().seen;
            }
        }
        reader.finish().unwrap();
        reader.into_sink().seen
    }

    #[test]
    fn the_reader_hands_every_event_to_the_sink() {
        assert_eq!(
            read_all(&[b"data: a\n\ndata: b\n", b"\ndata: c\n\n"]),
            vec!["a", "b", "c"]
        );
    }

    // Both adapters had this: the last event of a stream closed without a trailing newline was
    // only delivered because each of them remembered to drain the decoder afterwards.
    #[test]
    fn an_unterminated_final_event_still_reaches_the_sink() {
        assert_eq!(read_all(&[b"data: a\ndata: b"]), vec!["a", "b"]);
    }

    // The Anthropic adapter stops mid-stream when the UI channel closes; nothing after that point
    // may be delivered, including from the same chunk.
    #[test]
    fn a_sink_that_stops_sees_nothing_further() {
        assert_eq!(
            read_all(&[b"data: a\ndata: halt\ndata: c\n"]),
            vec!["a", "halt"]
        );
    }

    #[test]
    fn a_stream_cut_mid_character_is_recorded() {
        let _guard = crate::diag::test_lock();
        crate::diag::drain();

        // The first two bytes of `—` (E2 80 94) with the third never arriving.
        let mut decoder = SseDecoder::new();
        assert!(decoder.push_bytes(b"data: a\xe2\x80").is_empty());
        assert_eq!(decoder.flush(), Some("a\u{FFFD}".to_string()));

        let warnings = crate::diag::drain();
        assert!(
            warnings.iter().any(|w| w.contains("mid-character")),
            "expected a recorded warning, got {:?}",
            warnings
        );
    }
}