Fix websocket server functions and add an example

Cargo.toml

@@ -109,6 +109,7 @@ members = [
     "examples/fullstack-streaming",
     "examples/fullstack-desktop",
     "examples/fullstack-auth",
+    "examples/fullstack-websockets",
 
     # Playwright tests
     "packages/playwright-tests/liveview",

examples/fullstack-websockets/.gitignore

@@ -0,0 +1,4 @@
+dist
+target
+static
+.dioxus

examples/fullstack-websockets/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "fullstack-websocket-example"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+dioxus = { workspace = true, features = ["fullstack"] }
+futures.workspace = true
+tokio = { workspace = true, features = ["full"], optional = true }
+
+[features]
+server = ["dioxus/server", "dep:tokio"]
+web = ["dioxus/web"]

examples/fullstack-websockets/src/main.rs

@@ -0,0 +1,69 @@
+#![allow(non_snake_case)]
+use dioxus::prelude::{
+    server_fn::{codec::JsonEncoding, BoxedStream, Websocket},
+    *,
+};
+use futures::{channel::mpsc, SinkExt, StreamExt};
+
+fn main() {
+    launch(app);
+}
+
+fn app() -> Element {
+    let mut uppercase = use_signal(String::new);
+    let mut uppercase_channel = use_signal(|| None);
+
+    // Start the websocket connection in a background task
+    use_future(move || async move {
+        let (tx, rx) = mpsc::channel(1);
+        let mut receiver = uppercase_ws(rx.into()).await.unwrap();
+        // Store the channel in a signal for use in the input handler
+        uppercase_channel.set(Some(tx));
+        // Whenever we get a message from the server, update the uppercase signal
+        while let Some(Ok(msg)) = receiver.next().await {
+            uppercase.set(msg);
+        }
+    });
+
+    rsx! {
+        input {
+            oninput: move |e| async move {
+                if let Some(mut uppercase_channel) = uppercase_channel() {
+                    let msg = e.value();
+                    uppercase_channel.send(Ok(msg)).await.unwrap();
+                }
+            },
+        }
+        "Uppercase: {uppercase}"
+    }
+}
+
+// The server macro accepts a protocol parameter which implements the protocol trait. The protocol
+// controls how the inputs and outputs are encoded when handling the server function. In this case,
+// the websocket<json, json> protocol can encode a stream input and stream output where messages are
+// serialized as JSON
+#[server(protocol = Websocket<JsonEncoding, JsonEncoding>)]
+async fn uppercase_ws(
+    input: BoxedStream<String, ServerFnError>,
+) -> Result<BoxedStream<String, ServerFnError>, ServerFnError> {
+    let mut input = input;
+
+    // Create a channel with the output of the websocket
+    let (mut tx, rx) = mpsc::channel(1);
+
+    // Spawn a task that processes the input stream and sends any new messages to the output
+    tokio::spawn(async move {
+        while let Some(msg) = input.next().await {
+            if tx
+                .send(msg.map(|msg| msg.to_ascii_uppercase()))
+                .await
+                .is_err()
+            {
+                break;
+            }
+        }
+    });
+
+    // Return the output stream
+    Ok(rx.into())
+}

packages/playwright-tests/fullstack.spec.js

@@ -68,3 +68,11 @@ test("document elements", async ({ page }) => {
   const main = page.locator("#main");
   await expect(main).toHaveCSS("font-family", "Roboto");
 });
+
+test("websockets", async ({ page }) => {
+  await page.goto("http://localhost:3333");
+  // wait until the websocket div is mounted
+  const wsDiv = page.locator("div#websocket-div");
+  await expect(wsDiv).toHaveText("Received: HELLO WORLD");
+});
+

packages/playwright-tests/fullstack/Cargo.toml

@@ -8,6 +8,7 @@ publish = false
 
 [dependencies]
 dioxus = { workspace = true, features = ["fullstack"] }
+futures.workspace = true
 serde = "1.0.218"
 tokio = { workspace = true, features = ["full"], optional = true }
 

packages/playwright-tests/fullstack/src/main.rs

@@ -5,7 +5,14 @@
 // - Hydration
 
 #![allow(non_snake_case)]
-use dioxus::{prelude::*, CapturedError};
+use dioxus::{
+    prelude::{
+        server_fn::{codec::JsonEncoding, BoxedStream, Websocket},
+        *,
+    },
+    CapturedError,
+};
+use futures::{channel::mpsc, SinkExt, StreamExt};
 
 fn main() {
     dioxus::LaunchBuilder::new()
@@ -43,6 +50,7 @@ fn app() -> Element {
         OnMounted {}
         DefaultServerFnCodec {}
         DocumentElements {}
+        WebSockets {}
     }
 }
 
@@ -156,3 +164,42 @@ fn DocumentElements() -> Element {
         document::Style { id: "style-head", "body {{ font-family: 'Roboto'; }}" }
     }
 }
+
+#[server(protocol = Websocket<JsonEncoding, JsonEncoding>)]
+async fn echo_ws(
+    input: BoxedStream<String, ServerFnError>,
+) -> Result<BoxedStream<String, ServerFnError>, ServerFnError> {
+    let mut input = input;
+
+    let (mut tx, rx) = mpsc::channel(1);
+
+    tokio::spawn(async move {
+        while let Some(msg) = input.next().await {
+            let _ = tx.send(msg.map(|msg| msg.to_ascii_uppercase())).await;
+        }
+    });
+
+    Ok(rx.into())
+}
+
+/// This component tests websocket server functions
+#[component]
+fn WebSockets() -> Element {
+    let mut received = use_signal(String::new);
+    use_future(move || async move {
+        let (mut tx, rx) = mpsc::channel(1);
+        let mut receiver = echo_ws(rx.into()).await.unwrap();
+        tx.send(Ok("hello world".to_string())).await.unwrap();
+        while let Some(Ok(msg)) = receiver.next().await {
+            println!("Received: {}", msg);
+            received.set(msg);
+        }
+    });
+
+    rsx! {
+        div {
+            id: "websocket-div",
+            "Received: {received}"
+        }
+    }
+}

packages/server/Cargo.toml

@@ -42,6 +42,7 @@ tracing-futures = { workspace = true }
 once_cell = { workspace = true }
 async-trait = { workspace = true }
 serde = { workspace = true }
+enumset = "1.1.5"
 
 futures-util = { workspace = true }
 futures-channel = { workspace = true }