Websocket route affected by serve_dioxus_application in fullstack mode

[zak-grivell]

Not sure if I am just being silly but i cannot find a solution to this. I am trying to use web sockets within dioxus full stack,

When

.serve_dioxus_application(ServeConfigBuilder::default(), App)

is added to router the web socket runs and echo's correctly but obviously the html is not served whereas with it the websocket connects but then immediately disconnects when i run

wscat -c ws://172.16.4.60:59445/ws

i get both

Connected (press CTRL+C to quit)
Disconnected (code: 1006, reason: "")

immediately.

This is my Cargo.toml file

[package]
name = "quizpop"
version = "0.1.0"
authors = ["Zak Grivell <[email protected]>"]
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
axum = { version = "0.7.9", features = ["ws"], optional=true }
dioxus = { version = "0.6.0", features = ["router", "fullstack"] }
tokio = { version = "1.42.0", features = ["full"], optional = true }
dioxus-cli-config = { version = "*", optional = true }

[features]
default = []
# The feature that are only required for the web = ["dioxus/web"] build target should be optional and only enabled in the web = ["dioxus/web"] feature
web = ["dioxus/web"]
# The feature that are only required for the desktop = ["dioxus/desktop"] build target should be optional and only enabled in the desktop = ["dioxus/desktop"] feature
desktop = ["dioxus/desktop"]
# The feature that are only required for the mobile = ["dioxus/mobile"] build target should be optional and only enabled in the mobile = ["dioxus/mobile"] feature
mobile = ["dioxus/mobile"]

server = ["dioxus/server", "axum", "tokio", "dioxus-cli-config"]

[profile]

[profile.wasm-dev]
inherits = "dev"
opt-level = 1

[profile.server-dev]
inherits = "dev"

[profile.android-dev]
inherits = "dev"

this is my main.rs

use dioxus::prelude::*;

#[cfg(feature = "server")]
mod server {
    use super::App;
    use axum::{
        extract::ws::{Message, WebSocket, WebSocketUpgrade},
        response::Response,
        routing::get,
        Router,
    };
    use dioxus::prelude::DioxusRouterExt;
    use dioxus::prelude::*;
    use tokio::net::TcpListener;

    pub async fn server_init() {
        let app = Router::new()
            .route("/ws", get(websocket_handler))
            .serve_dioxus_application(ServeConfigBuilder::default(), App);

        let listener = TcpListener::bind(dioxus_cli_config::fullstack_address_or_localhost())
            .await
            .unwrap();

        axum::serve(listener, app).await.unwrap();
    }

    async fn websocket_handler(ws: WebSocketUpgrade) -> Response {
        ws.on_upgrade(handle_socket)
    }

    async fn handle_socket(mut socket: WebSocket) {
        if let Err(e) = socket
            .send(Message::Text("Hello from Rust WebSocket!".to_string()))
            .await
        {
            return;
        }

        loop {
            match socket.recv().await {
                Some(Ok(Message::Text(text))) => {
                    let response = format!("Server received: {}", text);

                    if let Err(e) = socket.send(Message::Text(response)).await {
                        break;
                    }
                }
                Some(Ok(Message::Binary(data))) => {
                    if let Err(e) = socket.send(Message::Binary(data)).await {
                        break;
                    }
                }
                Some(Ok(Message::Close(_))) => {
                    break;
                }
                Some(Ok(Message::Ping(data))) => {
                    if let Err(e) = socket.send(Message::Pong(data)).await {
                        break;
                    }
                }
                Some(Ok(Message::Pong(_))) => {}
                Some(Err(e)) => {
                    break;
                }
                None => {
                    break;
                }
            }
        }

        println!("WebSocket connection handler finished");
    }
}

#[cfg(feature = "server")]
#[tokio::main]
async fn main() {
    server::server_init().await;
}

#[component]
pub fn App() -> Element {
    rsx! { "HotDog!" }
}

#[cfg(not(feature = "server"))]
fn main() {
    dioxus::launch(App);
}

Any help would be greatly appreciated!