feat(rust): add standalone framing crate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-rust/Cargo.lock

@@ -268,6 +268,17 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "aws-smithy-eventstream"
+version = "0.61.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
+dependencies = [
+ "aws-smithy-types",
+ "bytes",
+ "crc32fast",
+]
+
 [[package]]
 name = "aws-smithy-http"
 version = "0.64.0"
@@ -1970,6 +1981,20 @@ dependencies = [
  "veil",
 ]
 
+[[package]]
+name = "litellm-framing"
+version = "0.1.0"
+dependencies = [
+ "aws-smithy-eventstream",
+ "aws-smithy-types",
+ "bytes",
+ "futures-util",
+ "rstest",
+ "sse-stream",
+ "thiserror 2.0.19",
+ "tokio",
+]
+
 [[package]]
 name = "litellm-python-bridge"
 version = "0.1.0"
@@ -3282,6 +3307,19 @@ dependencies = [
  "unicode-segmentation",
 ]
 
+[[package]]
+name = "sse-stream"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
+dependencies = [
+ "bytes",
+ "futures-util",
+ "http-body 1.1.0",
+ "http-body-util",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "stable_deref_trait"
 version = "1.2.1"

litellm-rust/crates/framer/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "litellm-framing"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[features]
+default = ["aws", "sse"]
+aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
+sse = ["dep:sse-stream"]
+
+[dependencies]
+aws-smithy-eventstream = { version = "=0.61.1", optional = true }
+aws-smithy-types = { version = "1.6.1", optional = true }
+bytes = "1"
+futures-util.workspace = true
+sse-stream = { version = "=0.2.6", optional = true }
+thiserror.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true
+tokio.workspace = true

litellm-rust/crates/framer/src/aws_event_stream.rs

@@ -0,0 +1,66 @@
+use bytes::{Buf, Bytes, BytesMut};
+use futures_util::{Stream, StreamExt};
+
+use aws_smithy_eventstream::frame::read_message_from;
+use aws_smithy_types::event_stream::Header;
+
+use crate::{Error, Framer};
+
+const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct AwsEventStreamFrame {
+    pub headers: Vec<Header>,
+    pub payload: Bytes,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct AwsEventStreamFramer;
+
+impl Framer for AwsEventStreamFramer {
+    type Frame = AwsEventStreamFrame;
+
+    fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
+    where
+        S: Stream<Item = Result<B, E>> + Send,
+        B: Buf + Send,
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        futures_util::stream::try_unfold(
+            (Box::pin(input), BytesMut::new()),
+            |(mut input, mut buffer)| async move {
+                loop {
+                    if buffer.len() >= 4 {
+                        let length = (&buffer[..4]).get_u32() as usize;
+                        if !(16..=MAX_FRAME_BYTES).contains(&length) {
+                            return Err(Error::InvalidLength(length));
+                        }
+                        if buffer.len() >= length {
+                            let raw = buffer.split_to(length).freeze();
+                            let message = read_message_from(raw)?;
+                            let frame = AwsEventStreamFrame {
+                                headers: message.headers().to_vec(),
+                                payload: message.payload().clone(),
+                            };
+                            return Ok(Some((frame, (input, buffer))));
+                        }
+                    }
+                    match input.next().await {
+                        Some(Ok(mut chunk)) => {
+                            while chunk.has_remaining() {
+                                let bytes = chunk.chunk();
+                                buffer.extend_from_slice(bytes);
+                                let length = bytes.len();
+                                chunk.advance(length);
+                            }
+                        }
+                        Some(Err(error)) => return Err(Error::Body(Box::new(error))),
+                        None if buffer.is_empty() => return Ok(None),
+                        None => return Err(Error::Truncated),
+                    }
+                }
+            },
+        )
+        .fuse()
+    }
+}

litellm-rust/crates/framer/src/error.rs

@@ -0,0 +1,17 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[cfg(feature = "sse")]
+    #[error("SSE framing failed: {0}")]
+    Sse(#[from] sse_stream::Error),
+    #[cfg(feature = "aws")]
+    #[error("AWS EventStream framing failed: {0}")]
+    Aws(#[from] aws_smithy_eventstream::error::Error),
+    #[error("body stream failed: {0}")]
+    Body(#[source] Box<dyn std::error::Error + Send + Sync>),
+    #[cfg(feature = "aws")]
+    #[error("invalid AWS EventStream frame length: {0}")]
+    InvalidLength(usize),
+    #[cfg(feature = "aws")]
+    #[error("truncated AWS EventStream frame")]
+    Truncated,
+}

litellm-rust/crates/framer/src/framer.rs

@@ -0,0 +1,13 @@
+use futures_util::Stream;
+
+use crate::Error;
+
+pub trait Framer: Send {
+    type Frame: Send;
+
+    fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
+    where
+        S: Stream<Item = Result<B, E>> + Send,
+        B: bytes::Buf + Send,
+        E: std::error::Error + Send + Sync + 'static;
+}

litellm-rust/crates/framer/src/lib.rs

@@ -0,0 +1,10 @@
+mod error;
+mod framer;
+
+pub use error::*;
+pub use framer::*;
+
+#[cfg(feature = "aws")]
+pub mod aws_event_stream;
+#[cfg(feature = "sse")]
+pub mod sse;

litellm-rust/crates/framer/src/sse.rs

@@ -0,0 +1,43 @@
+use futures_util::{Stream, StreamExt};
+
+use crate::{Error, Framer};
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct SseFrame {
+    pub event: Option<String>,
+    pub data: Option<String>,
+    pub id: Option<String>,
+    pub retry: Option<u64>,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct SseFramer;
+
+impl Framer for SseFramer {
+    type Frame = SseFrame;
+
+    fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<SseFrame, Error>> + Send
+    where
+        S: Stream<Item = Result<B, E>> + Send,
+        B: bytes::Buf + Send,
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input));
+        futures_util::stream::try_unfold(frames, |mut frames| async move {
+            let Some(frame) = frames.next().await else {
+                return Ok(None);
+            };
+            let frame = frame?;
+            Ok(Some((
+                SseFrame {
+                    event: frame.event,
+                    data: frame.data,
+                    id: frame.id,
+                    retry: frame.retry,
+                },
+                frames,
+            )))
+        })
+        .fuse()
+    }
+}

litellm-rust/crates/framer/tests/aws_event_stream.rs

@@ -0,0 +1,92 @@
+#![cfg(feature = "aws")]
+
+mod support;
+
+use std::io;
+
+use futures_util::TryStreamExt;
+use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
+use litellm_framing::{Error, Framer};
+use rstest::{fixture, rstest};
+
+use support::encode;
+
+async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result<Vec<AwsEventStreamFrame>, Error> {
+    AwsEventStreamFramer
+        .frame(futures_util::stream::iter(
+            bytes.chunks(chunk_size).map(Ok::<_, io::Error>),
+        ))
+        .try_collect()
+        .await
+}
+
+#[fixture]
+fn two_frames() -> Vec<u8> {
+    [encode(b"\xff\x00"), encode(b"second")].concat()
+}
+
+#[fixture]
+fn payload_frame() -> Vec<u8> {
+    encode(b"payload")
+}
+
+#[rstest]
+#[case(1)]
+#[case(3)]
+#[case(12)]
+#[case(usize::MAX)]
+#[tokio::test]
+async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads(
+    two_frames: Vec<u8>,
+    #[case] chunk_size: usize,
+) {
+    let chunk_size = chunk_size.min(two_frames.len());
+    let frames = collect_aws(&two_frames, chunk_size).await.unwrap();
+    assert_eq!(frames.len(), 2);
+    assert_eq!(frames[0].payload, &b"\xff\x00"[..]);
+    assert_eq!(frames[1].payload, "second");
+    assert_eq!(
+        frames[0].headers[0].value().as_string().unwrap().as_str(),
+        "payload"
+    );
+    assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7));
+}
+
+#[rstest]
+#[case(8)]
+#[case(0)]
+#[tokio::test]
+async fn rejects_corrupt_crcs(payload_frame: Vec<u8>, #[case] index: usize) {
+    let corrupt_index = if index == 0 {
+        payload_frame.len() - 1
+    } else {
+        index
+    };
+    let mut corrupt = payload_frame;
+    corrupt[corrupt_index] ^= 1;
+    assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_))));
+}
+
+#[rstest]
+#[case(0_u32)]
+#[case(15)]
+#[case(u32::MAX)]
+#[tokio::test]
+async fn rejects_invalid_lengths(#[case] length: u32) {
+    assert!(matches!(
+        collect_aws(&length.to_be_bytes(), 1).await,
+        Err(Error::InvalidLength(_))
+    ));
+}
+
+#[rstest]
+#[case(1)]
+#[case(3)]
+#[case(5)]
+#[tokio::test]
+async fn rejects_truncation(payload_frame: Vec<u8>, #[case] end: usize) {
+    assert!(matches!(
+        collect_aws(&payload_frame[..end], 1).await,
+        Err(Error::Truncated)
+    ));
+}

litellm-rust/crates/framer/tests/chaining.rs

@@ -0,0 +1,29 @@
+#![cfg(all(feature = "aws", feature = "sse"))]
+
+mod support;
+
+use std::io;
+
+use futures_util::TryStreamExt;
+use litellm_framing::Framer;
+use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
+use litellm_framing::sse::SseFramer;
+
+use support::encode;
+
+#[tokio::test]
+async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() {
+    let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat();
+    let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter(
+        bytes.chunks(3).map(Ok::<_, io::Error>),
+    ));
+    let frames = SseFramer
+        .frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload))
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+    assert_eq!(frames.len(), 1);
+    assert_eq!(frames[0].event.as_deref(), Some("delta"));
+    assert_eq!(frames[0].data.as_deref(), Some("hello"));
+    assert_eq!(frames[0].id.as_deref(), Some("7"));
+}

litellm-rust/crates/framer/tests/sse.rs

@@ -0,0 +1,67 @@
+#![cfg(feature = "sse")]
+
+use std::io;
+
+use futures_util::{StreamExt, TryStreamExt};
+use litellm_framing::sse::{SseFrame, SseFramer};
+use litellm_framing::{Error, Framer};
+use rstest::rstest;
+
+async fn collect_sse(chunks: &[&[u8]]) -> Result<Vec<SseFrame>, Error> {
+    SseFramer
+        .frame(futures_util::stream::iter(
+            chunks.iter().copied().map(Ok::<_, io::Error>),
+        ))
+        .try_collect()
+        .await
+}
+
+#[rstest]
+#[case(
+    &[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]],
+    vec![
+        SseFrame {
+            event: Some("delta".into()),
+            data: Some("€\nnext".into()),
+            id: Some("7".into()),
+            retry: Some(10),
+        },
+        SseFrame {
+            event: None,
+            data: Some("[DONE]".into()),
+            id: None,
+            retry: None,
+        },
+    ]
+)]
+#[tokio::test]
+async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel(
+    #[case] chunks: &[&[u8]],
+    #[case] expected: Vec<SseFrame>,
+) {
+    assert_eq!(collect_sse(chunks).await.unwrap(), expected);
+}
+
+#[tokio::test]
+async fn eof_does_not_dispatch_an_unterminated_frame() {
+    assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty());
+}
+
+#[rstest]
+#[case(io::ErrorKind::ConnectionReset)]
+#[case(io::ErrorKind::UnexpectedEof)]
+#[tokio::test]
+async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) {
+    let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([
+        Err(io::Error::new(kind, "reset")),
+        Ok(&b"data: later\n\n"[..]),
+    ])));
+    let error = frames.next().await.unwrap().unwrap_err();
+    assert!(matches!(
+        error,
+        Error::Sse(sse_stream::Error::Body(ref cause))
+            if cause.downcast_ref::<io::Error>().unwrap().kind() == kind
+    ));
+    assert!(frames.next().await.is_none());
+    assert!(frames.next().await.is_none());
+}

litellm-rust/crates/framer/tests/support/mod.rs

@@ -0,0 +1,15 @@
+use aws_smithy_eventstream::frame::write_message_to;
+use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
+use bytes::Bytes;
+
+pub fn encode(payload: &'static [u8]) -> Vec<u8> {
+    let message = Message::new(Bytes::from_static(payload))
+        .add_header(Header::new(
+            ":event-type",
+            HeaderValue::String("payload".into()),
+        ))
+        .add_header(Header::new("sequence", HeaderValue::Int32(7)));
+    let mut bytes = Vec::new();
+    write_message_to(&message, &mut bytes).unwrap();
+    bytes
+}