perf(anthropic): decode upstream SSE via shared SseDecoder

Replace the hand-rolled byte buffer + `find_double_newline` +
`parse_sse_frame` in `consume_and_forward` with the shared
`common::sse::SseDecoder`. The decoder provides cursor-tracked buffering
(single compact per batch instead of per-frame `Vec::drain`), deferred
UTF-8 validation for multi-byte characters split across network chunks,
and the same 1 MB DoS cap.

A small `resolve_event` helper preserves Anthropic's existing behavior of
inferring the event type from the payload's `"type"` field when an
explicit `event:` line is absent. Frame parsing is now zero-allocation
for the common single-line `data:` case. Tests updated to exercise the
decoder path; the inference fallback is covered explicitly.

Signed-off-by: XinyueZhang369 <zoeyzhang369@gmail.com>

Author: XinyueZhang369

model_gateway/src/routers/anthropic/sse.rs

@@ -3,7 +3,7 @@
 //! Provides SSE frame parsing, event formatting, stream wrappers,
 //! and the core stream consumption logic used by the streaming processor.
 
-use std::io;
+use std::{borrow::Cow, io};
 
 use axum::{
     body::Body,
@@ -18,7 +18,10 @@ use tokio::sync::mpsc;
 use tracing::{debug, error, warn};
 
 use super::mcp::{IterationResult, McpToolCall};
-use crate::routers::error::internal_error;
+use crate::routers::{
+    common::sse::{SseDecodeError, SseDecoder, SseFrame},
+    error::internal_error,
+};
 
 // ============================================================================
 // Constants
@@ -216,47 +219,43 @@ where
     F: Fn(&str) -> String,
 {
     let mut stream = response.bytes_stream();
-    let mut buffer = Vec::<u8>::new();
+    // Shared SSE decoder: cursor-tracked buffering, deferred UTF-8 validation
+    // (handles multi-byte characters split across network chunks), and built-in
+    // DoS protection at `MAX_SSE_BUFFER_SIZE`.
+    let mut decoder = SseDecoder::with_max_size(MAX_SSE_BUFFER_SIZE);
     let mut processor =
         EventProcessor::new(tx, global_index, is_first_iteration, resolve_server_name);
 
     while let Some(chunk_result) = stream.next().await {
         let chunk = chunk_result.map_err(|e| format!("Stream read error: {e}"))?;
 
-        // Guard against unbounded buffer growth (DoS protection).
-        // Check *before* extending so a single oversized chunk never
-        // causes an allocation beyond the cap.
-        if buffer.len() + chunk.len() > MAX_SSE_BUFFER_SIZE {
-            return Err(format!(
+        decoder.push(&chunk).map_err(|e| match e {
+            SseDecodeError::BufferOverflow => format!(
                 "SSE buffer exceeded maximum size ({MAX_SSE_BUFFER_SIZE} bytes) — possible malformed upstream stream"
-            ));
-        }
-
-        buffer.extend_from_slice(&chunk);
+            ),
+            other => format!("SSE decode error: {other}"),
+        })?;
 
-        // Process complete SSE frames (delimited by double newline).
-        // UTF-8 validation is deferred to complete frames so that multi-byte
-        // characters split across network chunks don't cause spurious errors.
-        while let Some(pos) = find_double_newline(&buffer) {
-            let frame_bytes = &buffer[..pos];
-            let frame = std::str::from_utf8(frame_bytes)
-                .map_err(|e| format!("Invalid UTF-8 in SSE frame: {e}"))?;
-            if let Some((event_type, data)) = parse_sse_frame(frame) {
+        while let Some(frame) = decoder.next_frame() {
+            let frame = frame.map_err(|e| format!("Invalid UTF-8 in SSE frame: {e}"))?;
+            if let Some((event_type, data)) = resolve_event(frame) {
                 processor.process(&event_type, &data).await?;
             }
-            buffer.drain(..pos + 2);
         }
+        decoder.compact();
     }
 
-    // Process any remaining data in buffer
-    if !buffer.is_empty() {
-        let remaining = std::str::from_utf8(&buffer)
-            .map_err(|e| format!("Invalid UTF-8 in final SSE data: {e}"))?;
-        let trimmed = remaining.trim();
-        if !trimmed.is_empty() {
-            if let Some((event_type, data)) = parse_sse_frame(trimmed) {
-                processor.process(&event_type, &data).await?;
-            }
+    // Process any trailing data not terminated by a blank line.
+    if let Some(frame) = decoder.flush() {
+        let frame = frame.map_err(|e| match e {
+            SseDecodeError::InvalidUtf8(u) => format!("Invalid UTF-8 in final SSE data: {u}"),
+            // The loop above drains every complete frame, so `flush()` can't
+            // return `IncompleteFlush` here — handled for an accurate message
+            // if that ever changes.
+            other => format!("SSE decode error on flush: {other}"),
+        })?;
+        if let Some((event_type, data)) = resolve_event(frame) {
+            processor.process(&event_type, &data).await?;
         }
     }
 
@@ -653,57 +652,33 @@ where
 }
 
 // ============================================================================
-// SSE frame parsing
+// SSE frame resolution
 // ============================================================================
 
-/// Find the position of `\n\n` in a byte buffer.
-fn find_double_newline(buf: &[u8]) -> Option<usize> {
-    buf.windows(2).position(|w| w == b"\n\n")
-}
-
-/// Parse a raw SSE frame into `(event_type, data)`.
+/// Resolve a decoded [`SseFrame`] into Anthropic's `(event_type, data)` pair.
 ///
-/// SSE frames look like:
-/// ```text
-/// event: content_block_start
-/// data: {"type":"content_block_start",...}
-/// ```
-fn parse_sse_frame(frame: &str) -> Option<(String, String)> {
-    let mut event_type = String::new();
-    let mut data_lines = Vec::new();
-
-    for line in frame.lines() {
-        let line = line.trim();
-        if line.is_empty() {
-            continue;
-        }
-        if let Some(value) = line.strip_prefix("event:") {
-            event_type = value.trim().to_string();
-        } else if let Some(value) = line.strip_prefix("data:") {
-            data_lines.push(value.trim().to_string());
-        }
-    }
-
-    if data_lines.is_empty() {
-        return None;
-    }
-
-    let data = data_lines.join("\n");
-
-    // If no event type specified, try to infer from data
-    if event_type.is_empty() {
-        if let Ok(parsed) = serde_json::from_str::<Value>(&data) {
-            if let Some(t) = parsed.get("type").and_then(|v| v.as_str()) {
-                event_type = t.to_string();
-            }
+/// Anthropic always emits an explicit `event:` line, but as a fallback
+/// (preserving prior behavior) the event type is inferred from the payload's
+/// `"type"` field when the `event:` line is absent or empty. Returns `None`
+/// when no event type can be determined.
+///
+/// The event type is returned as a `Cow` so the common case (explicit `event:`
+/// line) borrows straight from the decoded frame with no allocation; only the
+/// rare inference fallback allocates.
+fn resolve_event(frame: SseFrame<'_>) -> Option<(Cow<'_, str>, Cow<'_, str>)> {
+    let event_type = match frame.event_type {
+        Some(e) if !e.is_empty() => e,
+        _ => {
+            let parsed: Value = serde_json::from_str(&frame.data).ok()?;
+            Cow::Owned(parsed.get("type")?.as_str()?.to_string())
         }
-    }
+    };
 
     if event_type.is_empty() {
         return None;
     }
 
-    Some((event_type, data))
+    Some((event_type, frame.data))
 }
 
 // ============================================================================
@@ -714,34 +689,79 @@ fn parse_sse_frame(frame: &str) -> Option<(String, String)> {
 mod tests {
     use super::*;
 
+    /// Decode `bytes` through the shared `SseDecoder` and resolve each frame
+    /// to `(event_type, data)` the way `consume_and_forward` does.
+    fn decode_events(bytes: &[u8]) -> Vec<(String, String)> {
+        let mut decoder = SseDecoder::new();
+        decoder.push(bytes).unwrap();
+        let mut out = Vec::new();
+        while let Some(frame) = decoder.next_frame() {
+            if let Some((event_type, data)) = resolve_event(frame.unwrap()) {
+                out.push((event_type.into_owned(), data.into_owned()));
+            }
+        }
+        out
+    }
+
     #[test]
-    fn test_parse_sse_frame_basic() {
-        let frame = "event: message_start\ndata: {\"type\":\"message_start\"}";
-        let (event_type, data) = parse_sse_frame(frame).unwrap();
-        assert_eq!(event_type, "message_start");
+    fn test_resolve_event_basic() {
+        let frame = SseFrame {
+            event_type: Some(Cow::Borrowed("message_start")),
+            data: Cow::Borrowed("{\"type\":\"message_start\"}"),
+        };
+        let (event_type, data) = resolve_event(frame).unwrap();
+        assert_eq!(event_type.as_ref(), "message_start");
         assert_eq!(data, "{\"type\":\"message_start\"}");
     }
 
     #[test]
-    fn test_parse_sse_frame_content_block() {
-        let frame = "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}";
-        let (event_type, data) = parse_sse_frame(frame).unwrap();
-        assert_eq!(event_type, "content_block_start");
-        let parsed: Value = serde_json::from_str(&data).unwrap();
-        assert_eq!(parsed["index"], 0);
+    fn test_resolve_event_no_event_type_infers() {
+        // No `event:` line -> infer the type from the payload's "type" field.
+        let frame = SseFrame {
+            event_type: None,
+            data: Cow::Borrowed("{\"type\":\"ping\"}"),
+        };
+        let (event_type, _data) = resolve_event(frame).unwrap();
+        assert_eq!(event_type.as_ref(), "ping");
+    }
+
+    #[test]
+    fn test_resolve_event_uninferable_is_none() {
+        // No event type and no "type" field in the payload -> skipped.
+        let frame = SseFrame {
+            event_type: None,
+            data: Cow::Borrowed("{\"foo\":1}"),
+        };
+        assert!(resolve_event(frame).is_none());
     }
 
     #[test]
-    fn test_parse_sse_frame_no_event_type_infers() {
-        let frame = "data: {\"type\":\"ping\"}";
-        let (event_type, _data) = parse_sse_frame(frame).unwrap();
-        assert_eq!(event_type, "ping");
+    fn test_decode_events_basic() {
+        let events = decode_events(
+            b"event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: ping\ndata: {\"type\":\"ping\"}\n\n",
+        );
+        assert_eq!(events.len(), 2);
+        assert_eq!(events[0].0, "message_start");
+        assert_eq!(events[1].0, "ping");
+    }
+
+    #[test]
+    fn test_decode_events_content_block() {
+        let events = decode_events(b"event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n");
+        assert_eq!(events.len(), 1);
+        assert_eq!(events[0].0, "content_block_start");
+        let parsed: Value = serde_json::from_str(&events[0].1).unwrap();
+        assert_eq!(parsed["index"], 0);
     }
 
     #[test]
-    fn test_parse_sse_frame_empty() {
-        assert!(parse_sse_frame("").is_none());
-        assert!(parse_sse_frame("event: foo").is_none());
+    fn test_decode_events_infers_event_type() {
+        // A `data:`-only frame (no `event:` line) infers its type from the payload.
+        let events = decode_events(b"data: {\"type\":\"ping\"}\n\n");
+        assert_eq!(
+            events,
+            vec![("ping".to_string(), "{\"type\":\"ping\"}".to_string())]
+        );
     }
 
     #[test]
@@ -755,10 +775,19 @@ mod tests {
     }
 
     #[test]
-    fn test_parse_sse_frame_with_extra_whitespace() {
-        let frame = "  event: content_block_delta  \n  data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{}\"}}  ";
-        let (event_type, data) = parse_sse_frame(frame).unwrap();
-        assert_eq!(event_type, "content_block_delta");
+    fn test_decode_events_split_across_chunks() {
+        // A multi-byte-safe decoder must reassemble frames split mid-stream.
+        let mut decoder = SseDecoder::new();
+        decoder
+            .push(b"event: content_block_delta\ndata: {\"type\":\"content_block_de")
+            .unwrap();
+        assert!(decoder.next_frame().is_none()); // incomplete
+        decoder
+            .push(b"lta\",\"index\":1,\"delta\":{\"partial_json\":\"{}\"}}\n\n")
+            .unwrap();
+        let frame = decoder.next_frame().unwrap().unwrap();
+        let (event_type, data) = resolve_event(frame).unwrap();
+        assert_eq!(event_type.as_ref(), "content_block_delta");
         let parsed: Value = serde_json::from_str(&data).unwrap();
         assert_eq!(parsed["index"], 1);
     }