feat: support binary WebSocket frames in message handlers

Closes #1332

Propagate binary WebSocket frames (opcode 0x2) through the entire
Rust→Python pipeline instead of silently echoing them back.

- Add WsPayload enum (Text/Binary) to replace String throughout the
  WebSocket channel, registry messages, and send methods
- Forward ws::Message::Binary through the channel so Python handlers
  receive bytes for binary frames and str for text frames
- Accept str|bytes in send/broadcast methods, emitting the correct
  frame type on the wire
- Add receive() returning str|bytes, fix receive_text/receive_bytes
  to validate frame type

Made-with: Cursor

Author: sansyrox

robyn/ws.py

@@ -33,8 +33,8 @@ def __init__(self, websocket_connector: WebSocketConnector, channel=None):
         self._connector = websocket_connector
         self._channel = channel
 
-    async def receive_text(self) -> str:
-        """Receive the next text message. Blocks until a message arrives.
+    async def receive(self) -> str | bytes:
+        """Receive the next message. Returns str for text frames, bytes for binary frames.
         Raises WebSocketDisconnect when the connection is closed."""
         if self._channel is None:
             raise WebSocketDisconnect(reason="No message channel available")
@@ -43,31 +43,46 @@ async def receive_text(self) -> str:
             raise WebSocketDisconnect()
         return result
 
+    async def receive_text(self) -> str:
+        """Receive the next text frame. Raises TypeError if a binary frame arrives.
+        Raises WebSocketDisconnect when the connection is closed."""
+        msg = await self.receive()
+        if not isinstance(msg, str):
+            raise TypeError(f"Expected text frame, got {type(msg).__name__}")
+        return msg
+
     async def receive_bytes(self) -> bytes:
-        """Receive binary data (decoded from text)."""
-        text = await self.receive_text()
-        return text.encode("utf-8")
+        """Receive the next binary frame. Raises TypeError if a text frame arrives.
+        Raises WebSocketDisconnect when the connection is closed."""
+        msg = await self.receive()
+        if not isinstance(msg, bytes):
+            raise TypeError(f"Expected binary frame, got {type(msg).__name__}")
+        return msg
 
     async def receive_json(self):
-        """Receive and decode JSON data.
+        """Receive and decode JSON data from a text frame.
         Raises WebSocketDisconnect when the connection is closed."""
         text = await self.receive_text()
         return orjson.loads(text)
 
+    async def send(self, data: str | bytes):
+        """Send data to this WebSocket client. str sends a text frame, bytes sends a binary frame."""
+        await self._connector.async_send_to(self._connector.id, data)
+
     async def send_text(self, data: str):
-        """Send text data to this WebSocket client."""
+        """Send a text frame to this WebSocket client."""
         await self._connector.async_send_to(self._connector.id, data)
 
     async def send_bytes(self, data: bytes):
-        """Send binary data (as text) to this WebSocket client."""
-        await self._connector.async_send_to(self._connector.id, data.decode("utf-8"))
+        """Send a binary frame to this WebSocket client."""
+        await self._connector.async_send_to(self._connector.id, data)
 
     async def send_json(self, data):
-        """Send JSON data to this WebSocket client."""
+        """Send JSON data as a text frame to this WebSocket client."""
         await self.send_text(orjson.dumps(data).decode())
 
-    async def broadcast(self, data: str):
-        """Broadcast text data to all connected WebSocket clients on this endpoint."""
+    async def broadcast(self, data: str | bytes):
+        """Broadcast data to all connected WebSocket clients on this endpoint."""
         await self._connector.async_broadcast(data)
 
     async def close(self):

src/executors/web_socket_executors.rs

@@ -4,7 +4,20 @@ use pyo3::prelude::*;
 use pyo3_async_runtimes::TaskLocals;
 
 use crate::types::function_info::FunctionInfo;
-use crate::websockets::WebSocketConnector;
+use crate::websockets::{WebSocketConnector, WsPayload};
+
+fn extract_ws_return(_py: Python, output: &Bound<'_, PyAny>) -> Option<WsPayload> {
+    if output.is_none() {
+        return None;
+    }
+    if let Ok(s) = output.extract::<String>() {
+        Some(WsPayload::Text(s))
+    } else if let Ok(b) = output.extract::<Vec<u8>>() {
+        Some(WsPayload::Binary(b))
+    } else {
+        None
+    }
+}
 
 pub fn execute_ws_function(
     function: &FunctionInfo,
@@ -23,25 +36,23 @@ pub fn execute_ws_function(
         });
         let f = async {
             let output = fut.await.unwrap();
-            Python::with_gil(|py| output.extract::<Option<String>>(py).unwrap())
+            Python::with_gil(|py| extract_ws_return(py, output.bind(py)))
         }
         .into_actor(ws)
-        .map(|res, _, ctx| {
-            if let Some(msg) = res {
-                ctx.text(msg);
-            }
+        .map(|res, _, ctx| match res {
+            Some(WsPayload::Text(s)) => ctx.text(s),
+            Some(WsPayload::Binary(b)) => ctx.binary(b),
+            None => {}
         });
         ctx.spawn(f);
     } else {
         Python::with_gil(|py| {
             let handler = function.handler.bind(py).downcast().unwrap();
-            if let Some(op) = handler
-                .call1((ws.clone(),))
-                .unwrap()
-                .extract::<Option<String>>()
-                .unwrap()
-            {
-                ctx.text(op);
+            let result = handler.call1((ws.clone(),)).unwrap();
+            match extract_ws_return(py, &result) {
+                Some(WsPayload::Text(s)) => ctx.text(s),
+                Some(WsPayload::Binary(b)) => ctx.binary(b),
+                None => {}
             }
         });
     }

src/websockets/mod.rs

@@ -3,7 +3,7 @@ pub mod registry;
 use crate::executors::web_socket_executors::execute_ws_function;
 use crate::types::function_info::FunctionInfo;
 use crate::types::multimap::QueryParams;
-use registry::{Close, SendMessageToAll, SendText};
+use registry::{Close, SendMessage, SendMessageToAll};
 
 use actix::prelude::*;
 use actix::{Actor, AsyncContext, StreamHandler};
@@ -13,6 +13,7 @@ use log::debug;
 use once_cell::sync::OnceCell;
 use parking_lot::RwLock;
 use pyo3::prelude::*;
+use pyo3::types::PyBytes;
 use pyo3::IntoPyObject;
 use pyo3_async_runtimes::TaskLocals;
 use std::sync::Arc;
@@ -24,24 +25,47 @@ use crate::runtime;
 use registry::{Register, WebSocketRegistry};
 use std::collections::HashMap;
 
+#[derive(Clone)]
+pub enum WsPayload {
+    Text(String),
+    Binary(Vec<u8>),
+}
+
+fn extract_payload(message: &Bound<'_, PyAny>) -> PyResult<WsPayload> {
+    if let Ok(s) = message.extract::<String>() {
+        Ok(WsPayload::Text(s))
+    } else if let Ok(b) = message.extract::<Vec<u8>>() {
+        Ok(WsPayload::Binary(b))
+    } else {
+        Err(pyo3::exceptions::PyTypeError::new_err(
+            "message must be str or bytes",
+        ))
+    }
+}
+
 /// A Rust-backed channel receiver exposed to Python.
 /// Python handlers call `await channel.receive()` to get the next message.
-/// Returns the message string, or None when the connection is closed.
+/// Returns str for text frames, bytes for binary frames, or None when closed.
 #[pyclass]
 pub struct WebSocketChannel {
-    receiver: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<Option<String>>>>,
+    receiver: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<Option<WsPayload>>>>,
 }
 
 #[pymethods]
 impl WebSocketChannel {
     /// Await the next message from the WebSocket.
-    /// Returns the message string, or None if the connection was closed.
+    /// Returns str for text frames, bytes for binary frames, or None if closed.
     fn receive<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let receiver = self.receiver.clone();
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             let mut rx = receiver.lock().await;
             match rx.recv().await {
-                Some(Some(msg)) => Ok(Some(msg)),
+                Some(Some(WsPayload::Text(s))) => Python::with_gil(|py| {
+                    Ok(Some(s.into_pyobject(py).unwrap().into_any().unbind()))
+                }),
+                Some(Some(WsPayload::Binary(b))) => {
+                    Python::with_gil(|py| Ok(Some(PyBytes::new(py, &b).into_any().unbind())))
+                }
                 Some(None) | None => Ok(None),
             }
         })
@@ -56,9 +80,7 @@ pub struct WebSocketConnector {
     pub task_locals: TaskLocals,
     pub registry_addr: Addr<WebSocketRegistry>,
     pub query_params: QueryParams,
-    /// Sender side of the message channel (stays in the Actix actor).
-    pub message_sender: Option<mpsc::UnboundedSender<Option<String>>>,
-    /// Receiver side exposed to Python via WebSocketChannel.
+    pub message_sender: Option<mpsc::UnboundedSender<Option<WsPayload>>>,
     pub message_channel: Option<Py<WebSocketChannel>>,
 }
 
@@ -73,7 +95,7 @@ impl Actor for WebSocketConnector {
             addr: addr.clone(),
         });
 
-        let (tx, rx) = mpsc::unbounded_channel::<Option<String>>();
+        let (tx, rx) = mpsc::unbounded_channel::<Option<WsPayload>>();
         self.message_sender = Some(tx);
         self.message_channel = Python::with_gil(|py| {
             Some(
@@ -94,9 +116,6 @@ impl Actor for WebSocketConnector {
     }
 
     fn stopped(&mut self, ctx: &mut Self::Context) {
-        // Drop the sender to close the channel.
-        // This causes any pending `channel.receive()` in Python to return None,
-        // which the WebSocketAdapter converts to WebSocketDisconnect.
         self.message_sender.take();
 
         let function = self.router.get("close").unwrap();
@@ -123,15 +142,21 @@ impl Clone for WebSocketConnector {
     }
 }
 
-impl Handler<SendText> for WebSocketConnector {
+impl Handler<SendMessage> for WebSocketConnector {
     type Result = ();
 
-    fn handle(&mut self, msg: SendText, ctx: &mut Self::Context) {
+    fn handle(&mut self, msg: SendMessage, ctx: &mut Self::Context) {
         if self.id == msg.recipient_id {
-            ctx.text(msg.message.clone());
-            if msg.message == "Connection closed" {
-                // Close the WebSocket connection
-                ctx.stop();
+            match &msg.payload {
+                WsPayload::Text(s) => {
+                    ctx.text(s.clone());
+                    if s == "Connection closed" {
+                        ctx.stop();
+                    }
+                }
+                WsPayload::Binary(b) => {
+                    ctx.binary(b.clone());
+                }
             }
         }
     }
@@ -151,14 +176,17 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WebSocketConnecto
             Ok(ws::Message::Text(text)) => {
                 debug!("Text message received {:?}", text);
                 if let Some(ref sender) = self.message_sender {
-                    let _ = sender.send(Some(text.to_string()));
+                    let _ = sender.send(Some(WsPayload::Text(text.to_string())));
+                }
+            }
+            Ok(ws::Message::Binary(bin)) => {
+                debug!("Binary message received ({} bytes)", bin.len());
+                if let Some(ref sender) = self.message_sender {
+                    let _ = sender.send(Some(WsPayload::Binary(bin.to_vec())));
                 }
             }
-            Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
             Ok(ws::Message::Close(_close_reason)) => {
                 debug!("Socket was closed");
-                // Drop sender to signal channel closure so receive() returns None.
-                // The close handler is called once from stopped().
                 self.message_sender.take();
                 ctx.stop();
             }
@@ -169,63 +197,69 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WebSocketConnecto
 
 #[pymethods]
 impl WebSocketConnector {
-    pub fn sync_send_to(&self, recipient_id: String, message: String) {
+    pub fn sync_send_to(&self, recipient_id: String, message: &Bound<'_, PyAny>) -> PyResult<()> {
+        let payload = extract_payload(message)?;
         let recipient_id = Uuid::parse_str(&recipient_id).unwrap();
 
-        match self.registry_addr.try_send(SendText {
-            message,
+        match self.registry_addr.try_send(SendMessage {
+            payload,
             sender_id: self.id,
             recipient_id,
         }) {
-            Ok(_) => println!("Message sent successfully"),
-            Err(e) => println!("Failed to send message: {}", e),
+            Ok(_) => debug!("Message sent successfully"),
+            Err(e) => debug!("Failed to send message: {}", e),
         }
+        Ok(())
     }
 
     pub fn async_send_to(
         &self,
         py: Python,
         recipient_id: String,
-        message: String,
+        message: &Bound<'_, PyAny>,
     ) -> PyResult<Py<PyAny>> {
+        let payload = extract_payload(message)?;
         let registry = self.registry_addr.clone();
         let recipient_id = Uuid::parse_str(&recipient_id).unwrap();
         let sender_id = self.id;
 
         let awaitable = runtime::future_into_py(py, async move {
-            match registry.try_send(SendText {
-                message,
+            match registry.try_send(SendMessage {
+                payload,
                 sender_id,
                 recipient_id,
             }) {
-                Ok(_) => println!("Message sent successfully"),
-                Err(e) => println!("Failed to send message: {}", e),
+                Ok(_) => debug!("Message sent successfully"),
+                Err(e) => debug!("Failed to send message: {}", e),
             }
             Ok(())
         })?;
 
         Ok(awaitable.into_pyobject(py)?.into_any().into())
     }
 
-    pub fn sync_broadcast(&self, message: String) {
+    pub fn sync_broadcast(&self, message: &Bound<'_, PyAny>) -> PyResult<()> {
+        let payload = extract_payload(message)?;
         let registry = self.registry_addr.clone();
         match registry.try_send(SendMessageToAll {
-            message,
+            payload,
             sender_id: self.id,
         }) {
-            Ok(_) => println!("Message sent successfully"),
-            Err(e) => println!("Failed to send message: {}", e),
+            Ok(_) => debug!("Message sent successfully"),
+            Err(e) => debug!("Failed to send message: {}", e),
         }
+        Ok(())
     }
 
-    pub fn async_broadcast(&self, py: Python, message: String) -> PyResult<Py<PyAny>> {
+    pub fn async_broadcast(&self, py: Python, message: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+        let payload = extract_payload(message)?;
         let registry = self.registry_addr.clone();
         let sender_id = self.id;
 
         let awaitable = runtime::future_into_py(py, async move {
-            match registry.try_send(SendMessageToAll { message, sender_id }) {
-                Ok(_) => println!("Message sent successfully"),
-                Err(e) => println!("Failed to send message: {}", e),
+            match registry.try_send(SendMessageToAll { payload, sender_id }) {
+                Ok(_) => debug!("Message sent successfully"),
+                Err(e) => debug!("Failed to send message: {}", e),
             }
             Ok(())
         })?;
@@ -247,7 +281,6 @@ impl WebSocketConnector {
         self.query_params.clone()
     }
 
-    /// Get the message channel for WebSocket handlers.
     #[getter]
     pub fn get_message_channel(&self, py: Python) -> Option<Py<WebSocketChannel>> {
         self.message_channel.as_ref().map(|c| c.clone_ref(py))