Leptos Axum Socket

Realtime pub/sub communication for Leptos + Axum applications.

Usage

// Define the key and message types
#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Debug)]
pub struct MyKey {
    pub bla: String,
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct MyMsg {
    pub awesome_msg: String,
}

// Implement the SocketMsg trait for MyMsg to link the key and message types
impl SocketMsg for MyMsg {
    type Key = MyKey;
    #[cfg(feature = "ssr")]
    type AppState = AppState;
}

#[component]
pub fn MyComponent() -> impl IntoView {
    let socket = expect_socket_context();

    // Subscribe to receive messages that are sent with the given key
    socket.subscribe(
        MyKey {
            bla: "bla".to_string(),
        },
        |msg: &MyMsg| {
            // Simply log the message
            leptos::logging::log!("message: {msg:#?}");
        },
    );

    let on_click = move || {
        // Send a message with the given key
        socket.send(
            MyKey {
                bla: "bla".to_string(),
            },
            MyMsg {
                awesome_msg: "awesome message".to_string(),
            },
        );
    };

    view! { "..." }
}

#[server]
pub async fn my_server_function() -> Result<(), ServerFnError> {
    // Send from the server
    leptos_axum_socket::send(
       &MyKey {
           bla: "bla".to_string(),
       },
       &MyMsg {
           awesome_msg: "Hello, world!".to_string(),
       },
    ).await;

    Ok(())
}

For this to work you have to prepare a little bit.

Define your app state in your lib.rs:

use leptos::prelude::*;

#[cfg(feature = "ssr")]
#[derive(Clone, axum::extract::FromRef)]
pub struct AppState {
    // This is required for Leptos Axum Socket to work
    pub socket: leptos_axum_socket::ServerSocket,

    // this is required for Leptos to work with axum
    pub leptos_options: LeptosOptions,
}

Initialize your Axum app (probably in main.rs):

// Construct the Axum app state
let state = AppState {
    leptos_options: conf.leptos_options,
    server_socket: ServerSocket::new(),
};

// Optional: add subscription filters and message mappers
{
    let mut server_socket = state.server_socket.lock().await;
    server_socket.add_subscribe_filter(|key: MyKey, _ctx: &()| { key.bla == "bla" });
    server_socket.add_send_mapper(|key: MyKey, msg: MyMsg, _ctx: &()| {
        if key.bla == "bla" {
            Some(MyMsg {
                awesome_msg: msg.awesome_msg.replace("old", "new"),
            })
        } else {
            None
        }
    });
}

// Init the Axum app
let app = Router::new()
    .leptos_routes(&state, routes, {
        use basic::app::shell;

        let leptos_options = state.leptos_options.clone();
        move || shell(leptos_options.clone())
    })
    .socket_route(connect_to_websocket)    // Register the socket route (implementation below)
    .fallback(leptos_axum::file_and_error_handler::<AppState, _>(shell))
    .with_state(state);    // Register the state

Implement the connect_to_websocket handler:

#[cfg(feature = "ssr")]
pub async fn connect_to_websocket(
    ws: WebSocketUpgrade,
    State(socket): State<ServerSocket>,
) -> Response {
    // You could do authentication here

    // Provide extra context like the user's ID for example that is passed to the permission filters
    let ctx = ();

    upgrade_websocket( ws, socket, ctx)
}

And finally provide the context in your root Leptos component:

#[component]
pub fn App() -> impl IntoView {
    provide_socket_context();

    view! { "..." }
}