vex_web/
lib.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4use axum::Router;
5use axum::extract::ws::{Message, WebSocket};
6use axum::extract::{State, WebSocketUpgrade};
7use axum::http::HeaderMap;
8use axum::response::{Html, IntoResponse, Response};
9use axum::routing::get;
10use futures_util::{SinkExt, StreamExt};
11use rust_embed::Embed;
12use tracing::info;
13use vex_hub::{FrontendCommand, FrontendEvent, Hub};
14
15#[derive(Embed)]
16#[folder = "src/static/"]
17struct Asset;
18
19struct AppState {
20    hub: Arc<Hub>,
21}
22
23pub async fn run(hub: Arc<Hub>, port: u16) -> Result<()> {
24    let state = Arc::new(AppState { hub });
25
26    let app = Router::new()
27        .route("/", get(index_handler))
28        .route("/ws", get(ws_handler))
29        .route("/static/{*path}", get(static_handler))
30        .fallback(get(index_handler))
31        .with_state(state);
32
33    let listener = tokio::net::TcpListener::bind(("0.0.0.0", port)).await?;
34    info!("web UI listening on http://0.0.0.0:{}", port);
35
36    axum::serve(listener, app).await?;
37    Ok(())
38}
39
40async fn index_handler() -> Html<String> {
41    match Asset::get("index.html") {
42        Some(content) => Html(String::from_utf8_lossy(&content.data).to_string()),
43        None => Html("<h1>vex web UI</h1><p>static files not found</p>".to_string()),
44    }
45}
46
47async fn static_handler(axum::extract::Path(path): axum::extract::Path<String>) -> Response {
48    match Asset::get(&path) {
49        Some(content) => {
50            let mime = mime_guess::from_path(&path).first_or_octet_stream();
51            let mut headers = HeaderMap::new();
52            headers.insert(
53                axum::http::header::CONTENT_TYPE,
54                mime.as_ref().parse().unwrap(),
55            );
56            if path.ends_with(".css") || path.ends_with(".js") {
57                headers.insert(
58                    axum::http::header::CACHE_CONTROL,
59                    "public, max-age=3600".parse().unwrap(),
60                );
61            }
62            (headers, content.data.to_vec()).into_response()
63        }
64        None => axum::http::StatusCode::NOT_FOUND.into_response(),
65    }
66}
67
68async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<AppState>>) -> Response {
69    ws.on_upgrade(move |socket| handle_ws(socket, state))
70}
71
72async fn handle_ws(socket: WebSocket, state: Arc<AppState>) {
73    let (mut ws_tx, mut ws_rx) = socket.split();
74    let mut event_rx = state.hub.event_rx();
75    let command_tx = state.hub.command_tx();
76
77    // Send initial state
78    let initial_state = state.hub.state_rx().borrow().clone();
79    let json = serde_json::json!({
80        "type": "state_snapshot",
81        "state": initial_state,
82    })
83    .to_string();
84    let _ = ws_tx.send(Message::Text(json.into())).await;
85
86    loop {
87        tokio::select! {
88            // Hub events -> WebSocket
89            Ok(event) = event_rx.recv() => {
90                let json = event_to_json(&event);
91                if ws_tx.send(Message::Text(json.into())).await.is_err() {
92                    break;
93                }
94            }
95            // WebSocket messages -> Hub commands
96            Some(Ok(msg)) = ws_rx.next() => {
97                match msg {
98                    Message::Text(text) => {
99                        if let Ok(cmd) = serde_json::from_str::<FrontendCommand>(&text) {
100                            let _ = command_tx.send(cmd).await;
101                        }
102                    }
103                    Message::Close(_) => break,
104                    _ => {}
105                }
106            }
107            else => break,
108        }
109    }
110}
111
112fn event_to_json(event: &FrontendEvent) -> String {
113    match event {
114        FrontendEvent::StateSnapshot(state) => serde_json::json!({
115            "type": "state_snapshot",
116            "state": state,
117        })
118        .to_string(),
119        FrontendEvent::ShellOutput { shell_id, data } => serde_json::json!({
120            "type": "shell_output",
121            "shell_id": shell_id.to_string(),
122            "data": data,
123        })
124        .to_string(),
125        FrontendEvent::AgentConversationLine { agent_id, line } => serde_json::json!({
126            "type": "agent_conversation_line",
127            "agent_id": agent_id,
128            "line": line,
129        })
130        .to_string(),
131        FrontendEvent::AgentWatchEnd { agent_id } => serde_json::json!({
132            "type": "agent_watch_end",
133            "agent_id": agent_id,
134        })
135        .to_string(),
136        FrontendEvent::ShellBell { shell_id } => serde_json::json!({
137            "type": "shell_bell",
138            "shell_id": shell_id.to_string(),
139        })
140        .to_string(),
141    }
142}
vex_web/lib.rs

vex_web/
lib.rs
