fix: sse downloads

Author: sansyrox

integration_tests/base_routes.py

@@ -7,7 +7,7 @@
 from typing import Optional
 
 from integration_tests.subroutes import di_subrouter, static_router, sub_router
-from robyn import Headers, Request, Response, Robyn, SSEMessage, SSEResponse, WebSocketDisconnect, jsonify, serve_file, serve_html
+from robyn import Headers, Request, Response, Robyn, SSEMessage, SSEResponse, StreamingResponse, WebSocketDisconnect, jsonify, serve_file, serve_html
 from robyn.authentication import AuthenticationHandler, BearerGetter, Identity
 from robyn.robyn import QueryParams, Url
 from robyn.templating import JinjaTemplate
@@ -1353,6 +1353,63 @@ def event_generator():
     return SSEResponse(event_generator(), status_code=201)
 
 
+# --- Binary streaming endpoints ---
+
+
+@app.get("/stream/bytes")
+def stream_bytes(request):
+    """Stream binary data using bytes chunks"""
+
+    def bytes_generator():
+        # Generate 3 chunks of known binary data
+        for i in range(3):
+            yield bytes([i] * 1024)  # 1KB chunks filled with the chunk index
+
+    return StreamingResponse(
+        content=bytes_generator(),
+        media_type="application/octet-stream",
+        headers=Headers({"Content-Type": "application/octet-stream"}),
+    )
+
+
+@app.get("/stream/bytes_file")
+def stream_bytes_file(request):
+    """Stream a file in binary mode using yield from"""
+    test_file = os.path.join(current_file_path, "build", "index.html")
+
+    def file_generator():
+        with open(test_file, "rb") as f:
+            while True:
+                chunk = f.read(512)
+                if not chunk:
+                    break
+                yield chunk
+
+    return StreamingResponse(
+        content=file_generator(),
+        media_type="application/octet-stream",
+        headers=Headers({
+            "Content-Type": "application/octet-stream",
+            "Content-Disposition": "attachment; filename=index.html",
+        }),
+    )
+
+
+@app.get("/stream/mixed_text")
+def stream_mixed_text(request):
+    """Stream text data using string chunks (ensures str still works)"""
+
+    def text_generator():
+        for i in range(3):
+            yield f"text chunk {i}\n"
+
+    return StreamingResponse(
+        content=text_generator(),
+        media_type="text/plain",
+        headers=Headers({"Content-Type": "text/plain"}),
+    )
+
+
 def main():
     app.set_response_header("server", "robyn")
     app.serve_directory(

integration_tests/test_binary_streaming.py

@@ -0,0 +1,81 @@
+import os
+
+import pytest
+import requests
+
+from integration_tests.helpers.http_methods_helpers import BASE_URL
+
+
+@pytest.mark.benchmark
+def test_stream_bytes_basic(session):
+    """Test that binary bytes can be streamed without error"""
+    response = requests.get(f"{BASE_URL}/stream/bytes", stream=True, timeout=5)
+    assert response.status_code == 200
+    assert response.headers.get("Content-Type") == "application/octet-stream"
+
+    # Collect all streamed data
+    data = b""
+    for chunk in response.iter_content(chunk_size=None):
+        if chunk:
+            data += chunk
+
+    # We expect 3 chunks of 1024 bytes each
+    assert len(data) == 3 * 1024
+
+    # Verify chunk contents: chunk i is filled with byte value i
+    for i in range(3):
+        chunk = data[i * 1024 : (i + 1) * 1024]
+        assert chunk == bytes([i] * 1024), f"Chunk {i} has unexpected content"
+
+
+@pytest.mark.benchmark
+def test_stream_bytes_no_sse_headers(session):
+    """Test that binary streaming responses do NOT include SSE-specific headers"""
+    response = requests.get(f"{BASE_URL}/stream/bytes", stream=True, timeout=5)
+    assert response.status_code == 200
+
+    # SSE-specific headers should NOT be present for binary streams
+    assert response.headers.get("X-Accel-Buffering") is None
+    assert response.headers.get("Pragma") is None
+    assert response.headers.get("Expires") is None
+
+
+@pytest.mark.benchmark
+def test_stream_bytes_file(session):
+    """Test streaming a file in binary mode"""
+    response = requests.get(f"{BASE_URL}/stream/bytes_file", stream=True, timeout=5)
+    assert response.status_code == 200
+    assert response.headers.get("Content-Type") == "application/octet-stream"
+    assert "attachment" in response.headers.get("Content-Disposition", "")
+
+    # Collect all streamed data
+    streamed_data = b""
+    for chunk in response.iter_content(chunk_size=None):
+        if chunk:
+            streamed_data += chunk
+
+    # Read the original file to compare
+    test_file = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "build", "index.html"
+    )
+    with open(test_file, "rb") as f:
+        original_data = f.read()
+
+    assert streamed_data == original_data, "Streamed file content does not match original"
+
+
+@pytest.mark.benchmark
+def test_stream_text_still_works(session):
+    """Test that string-based streaming still works after the bytes change"""
+    response = requests.get(f"{BASE_URL}/stream/mixed_text", stream=True, timeout=5)
+    assert response.status_code == 200
+    assert response.headers.get("Content-Type") == "text/plain"
+
+    content = b""
+    for chunk in response.iter_content(chunk_size=None):
+        if chunk:
+            content += chunk
+
+    text = content.decode("utf-8")
+    for i in range(3):
+        assert f"text chunk {i}" in text

robyn/responses.py

@@ -65,7 +65,7 @@ def serve_file(file_path: str, file_name: Optional[str] = None) -> FileResponse:
 class AsyncGeneratorWrapper:
     """Optimized true-streaming wrapper for async generators"""
 
-    def __init__(self, async_gen: AsyncGenerator[str, None]):
+    def __init__(self, async_gen: AsyncGenerator[Union[str, bytes], None]):
         self.async_gen = async_gen
         self._loop = None
         self._iterator = None
@@ -124,7 +124,10 @@ async def get_next():
 class StreamingResponse:
     def __init__(
         self,
-        content: Union[Generator[str, None, None], AsyncGenerator[str, None]],
+        content: Union[
+            Generator[Union[str, bytes], None, None],
+            AsyncGenerator[Union[str, bytes], None],
+        ],
         status_code: Optional[int] = None,
         headers: Optional[Headers] = None,
         media_type: str = "text/event-stream",
@@ -149,7 +152,10 @@ def __init__(
 
 
 def SSEResponse(
-    content: Union[Generator[str, None, None], AsyncGenerator[str, None]],
+    content: Union[
+        Generator[Union[str, bytes], None, None],
+        AsyncGenerator[Union[str, bytes], None],
+    ],
     status_code: Optional[int] = None,
     headers: Optional[Headers] = None,
 ) -> StreamingResponse:

src/types/response.rs

@@ -32,6 +32,7 @@ pub struct StreamingResponse {
     pub status_code: u16,
     pub headers: Headers,
     pub content_generator: Py<PyAny>,
+    pub media_type: String,
 }
 
 #[derive(Debug)]
@@ -86,11 +87,17 @@ impl Responder for Response {
 }
 
 impl StreamingResponse {
-    pub fn new(status_code: u16, headers: Headers, content_generator: Py<PyAny>) -> Self {
+    pub fn new(
+        status_code: u16,
+        headers: Headers,
+        content_generator: Py<PyAny>,
+        media_type: String,
+    ) -> Self {
         Self {
             status_code,
             headers,
             content_generator,
+            media_type,
         }
     }
 }
@@ -105,13 +112,15 @@ impl Responder for StreamingResponse {
 
         apply_hashmap_headers(&mut response_builder, &self.headers);
 
-        // Optimized headers for SSE streaming
-        response_builder
-            .append_header(("Connection", "keep-alive"))
-            .append_header(("X-Accel-Buffering", "no")) // Disable nginx buffering
-            .append_header(("Cache-Control", "no-cache, no-store, must-revalidate"))
-            .append_header(("Pragma", "no-cache"))
-            .append_header(("Expires", "0"));
+        // Only add SSE-specific headers for event-stream responses
+        if self.media_type == "text/event-stream" {
+            response_builder
+                .append_header(("Connection", "keep-alive"))
+                .append_header(("X-Accel-Buffering", "no")) // Disable nginx buffering
+                .append_header(("Cache-Control", "no-cache, no-store, must-revalidate"))
+                .append_header(("Pragma", "no-cache"))
+                .append_header(("Expires", "0"));
+        }
 
         // Create the optimized stream from the Python generator
         let stream = create_python_stream(self.content_generator);
@@ -142,18 +151,29 @@ fn create_python_stream(
                     // Try to get the next value from the generator (sync)
                     match gen.call_method0("__next__") {
                         Ok(value) => {
-                            match value.extract::<String>() {
-                                Ok(string_value) => {
-                                    debug!("Generator yielded: {}", string_value);
-                                    Some((string_value, generator))
-                                }
-                                Err(extract_err) => {
-                                    debug!(
-                                        "Failed to extract string from generator value: {}",
-                                        extract_err
-                                    );
-                                    None // End of stream
-                                }
+                            // Try bytes first (common for binary file streaming),
+                            // then fall back to string extraction
+                            if let Ok(py_bytes) = value.downcast::<PyBytes>() {
+                                let data = py_bytes.as_bytes().to_vec();
+                                debug!("Generator yielded {} bytes", data.len());
+                                Some((data, generator))
+                            } else if let Ok(string_value) = value.extract::<String>() {
+                                debug!(
+                                    "Generator yielded string of len {}",
+                                    string_value.len()
+                                );
+                                Some((string_value.into_bytes(), generator))
+                            } else {
+                                let type_name = value
+                                    .get_type()
+                                    .name()
+                                    .map(|n| n.to_string())
+                                    .unwrap_or_else(|_| "unknown".to_string());
+                                debug!(
+                                    "Generator yielded unsupported type: {}",
+                                    type_name
+                                );
+                                None // End of stream
                             }
                         }
                         Err(call_err) => {
@@ -171,7 +191,7 @@ fn create_python_stream(
         })
         .await
         {
-            Ok(Some((string_value, generator))) => Some((Ok(Bytes::from(string_value)), generator)),
+            Ok(Some((data, generator))) => Some((Ok(Bytes::from(data)), generator)),
             Ok(None) => None,
             Err(join_err) => {
                 debug!(
@@ -583,6 +603,6 @@ impl FromPyObject<'_, '_> for StreamingResponse {
             "Successfully extracted StreamingResponse with status {} from type {}",
             status_code, type_name
         );
-        Ok(StreamingResponse::new(status_code, headers, content))
+        Ok(StreamingResponse::new(status_code, headers, content, media_type))
     }
 }