hanzo-api 1.1.12

Http Api for Hanzo AI platform
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

use crate::api_ws::api_ws_handlers::ws_handler;
use hanzo_messages::hanzo_utils::hanzo_logging::hanzo_log;
use hanzo_messages::hanzo_utils::hanzo_logging::HanzoLogLevel;
use hanzo_messages::hanzo_utils::hanzo_logging::HanzoLogOption;
use warp::{http::StatusCode, Filter, Rejection, Reply};

/// Handle rejections from the routes
async fn handle_rejection(err: Rejection) -> Result<impl Reply, Rejection> {
    if err.is_not_found() {
        return Ok(warp::reply::with_status("Not Found", StatusCode::NOT_FOUND));
    }

    // Log the error
    hanzo_log(
        HanzoLogOption::WsAPI,
        HanzoLogLevel::Error,
        &format!("unhandled rejection: {:?}", err),
    );

    // Return a generic error message
    Ok(warp::reply::with_status(
        "Internal Server Error",
        StatusCode::INTERNAL_SERVER_ERROR,
    ))
}

/// Create the Warp routes for WebSocket endpoints
pub fn ws_routes(
    ws_address: std::net::SocketAddr,
) -> impl Filter<Extract = impl warp::Reply, Error = Rejection> + Clone {
    tracing::info!("Setting up WebSocket routes");

    let root_ws = warp::path::end().and(warp::ws()).map(move |ws| {
        let ws: warp::ws::Ws = ws;
        ws.on_upgrade(move |socket| ws_handler(socket, ws_address))
    });

    root_ws
        .with(warp::cors().allow_any_origin())
        .with(warp::log("websocket"))
        .recover(handle_rejection)
}