Skip to main content

webui_dev_server/
livereload.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT license.
3
4//! Server-Sent Events live-reload broadcaster.
5//!
6//! [`LiveReload`] holds a `tokio::sync::broadcast` channel and an actix
7//! handler that streams events to connected browsers. Cheap to clone — the
8//! underlying `Sender` is itself an `Arc`-like type.
9//!
10//! ## Wire protocol
11//!
12//! The browser opens an `EventSource` connection and listens for two
13//! named events:
14//!
15//! - `reload` — successful rebuild; the page should refresh.
16//! - `reload-error` — rebuild failed; the page logs the message and waits
17//!   for the next event without reloading.
18//!
19//! A 30-second `:heartbeat` comment keeps the connection alive through
20//! intermediate proxies that drop idle TCP streams.
21
22use std::time::Duration;
23
24use actix_web::http::header::{CACHE_CONTROL, CONTENT_TYPE};
25use actix_web::web::Bytes;
26use actix_web::{web, HttpResponse};
27use futures_util::stream::Stream;
28use tokio::sync::broadcast;
29
30use crate::inject::inject_before_body_close;
31
32/// Broadcast channel capacity. Live reload only emits one event per build
33/// — 16 is plenty for a dev server, and lagging subscribers drop old
34/// events rather than delaying the publisher.
35const RELOAD_CHANNEL_CAPACITY: usize = 16;
36
37/// Heartbeat interval for SSE streams so intermediaries don't drop idle
38/// connections.
39const SSE_HEARTBEAT: Duration = Duration::from_secs(30);
40
41/// A live-reload event. Keep this enum stable; both the wire format and
42/// the broadcast subscribers depend on it.
43#[derive(Clone, Debug)]
44pub enum ReloadEvent {
45    /// Successful rebuild — clients should reload.
46    Reload,
47    /// Build failed — clients log the message but do not reload.
48    Error(String),
49}
50
51/// SSE live-reload broadcaster. Cheap to clone; share across handlers.
52#[derive(Clone)]
53pub struct LiveReload {
54    endpoint: String,
55    tx: broadcast::Sender<ReloadEvent>,
56    client_script: String,
57}
58
59impl LiveReload {
60    /// Create a new broadcaster.
61    ///
62    /// `endpoint` should be a root-relative URL path (e.g.
63    /// `"/__webui/livereload"`). Use a root-relative path so the script
64    /// works when the served pages set `<base href>` for sub-path
65    /// deployments.
66    #[must_use]
67    pub fn new(endpoint: impl Into<String>) -> Self {
68        let endpoint = endpoint.into();
69        let (tx, _rx) = broadcast::channel::<ReloadEvent>(RELOAD_CHANNEL_CAPACITY);
70        let client_script = build_client_script(&endpoint);
71        Self {
72            endpoint,
73            tx,
74            client_script,
75        }
76    }
77
78    /// The endpoint URL where the SSE stream is served.
79    #[must_use]
80    pub fn endpoint(&self) -> &str {
81        &self.endpoint
82    }
83
84    /// The inline `<script>...</script>` that subscribes to this endpoint.
85    /// Inject this into served HTML.
86    #[must_use]
87    pub fn client_script(&self) -> &str {
88        &self.client_script
89    }
90
91    /// Inject [`Self::client_script`] immediately before `</body>` in
92    /// `html`. Appends to the end if no closing tag is found.
93    #[must_use]
94    pub fn inject(&self, html: &str) -> String {
95        inject_before_body_close(html, &self.client_script)
96    }
97
98    /// Broadcast a `Reload` event to all connected clients. A failed send
99    /// (no subscribers) is ignored — that's a normal state when no browser
100    /// is currently connected.
101    pub fn broadcast_reload(&self) {
102        let _ = self.tx.send(ReloadEvent::Reload);
103    }
104
105    /// Broadcast a `reload-error` event. Newlines in the message are
106    /// replaced with spaces to keep the SSE frame on one line.
107    pub fn broadcast_error(&self, msg: impl Into<String>) {
108        let _ = self.tx.send(ReloadEvent::Error(msg.into()));
109    }
110
111    /// Direct access to the underlying broadcaster, e.g. to count
112    /// subscribers (`sender().receiver_count()`).
113    #[must_use]
114    pub fn sender(&self) -> &broadcast::Sender<ReloadEvent> {
115        &self.tx
116    }
117}
118
119/// Actix handler for the SSE endpoint.
120///
121/// Wire it up in your `App` with:
122///
123/// ```ignore
124/// use actix_web::{web, App};
125/// use webui_dev_server::{livereload, LiveReload};
126///
127/// let lr = LiveReload::new("/__webui/livereload");
128/// let app = App::new()
129///     .app_data(web::Data::new(lr.clone()))
130///     .route(lr.endpoint(), web::get().to(livereload::sse_handler));
131/// ```
132pub async fn sse_handler(lr: web::Data<LiveReload>) -> HttpResponse {
133    let rx = lr.tx.subscribe();
134    HttpResponse::Ok()
135        .insert_header((CONTENT_TYPE, "text/event-stream"))
136        .insert_header((CACHE_CONTROL, "no-cache"))
137        // Disable nginx-style proxy buffering if anyone proxies dev.
138        .insert_header(("X-Accel-Buffering", "no"))
139        .streaming(sse_stream(rx))
140}
141
142fn sse_stream(
143    mut rx: broadcast::Receiver<ReloadEvent>,
144) -> impl Stream<Item = std::result::Result<Bytes, actix_web::Error>> {
145    async_stream::stream! {
146        // Hello frame so EventSource immediately marks the connection open.
147        yield Ok::<_, actix_web::Error>(Bytes::from_static(b": connected\n\n"));
148
149        let mut heartbeat = tokio::time::interval(SSE_HEARTBEAT);
150        heartbeat.tick().await; // skip immediate fire
151
152        loop {
153            tokio::select! {
154                _ = heartbeat.tick() => {
155                    yield Ok(Bytes::from_static(b": heartbeat\n\n"));
156                }
157                msg = rx.recv() => {
158                    match msg {
159                        Ok(ReloadEvent::Reload) => {
160                            yield Ok(Bytes::from_static(b"event: reload\ndata: ok\n\n"));
161                        }
162                        Ok(ReloadEvent::Error(e)) => {
163                            let payload = format!(
164                                "event: reload-error\ndata: {}\n\n",
165                                e.replace('\n', " ")
166                            );
167                            yield Ok(Bytes::from(payload));
168                        }
169                        Err(broadcast::error::RecvError::Lagged(_)) => continue,
170                        Err(broadcast::error::RecvError::Closed) => break,
171                    }
172                }
173            }
174        }
175    }
176}
177
178fn build_client_script(endpoint: &str) -> String {
179    // `EventSource` reconnects automatically; we only need to handle
180    // named events. The endpoint is embedded as a JSON-quoted string so
181    // any unusual characters survive (we don't expect them, but it's
182    // cheap insurance).
183    let endpoint_js = json_string_literal(endpoint);
184    format!(
185        "<script>(function(){{try{{var s=new EventSource({endpoint_js});\
186s.addEventListener(\"reload\",function(){{location.reload();}});\
187s.addEventListener(\"reload-error\",function(e){{console.error(\"[webui-dev] rebuild failed:\",e.data);}});\
188}}catch(e){{console.warn(\"[webui-dev] live reload unavailable:\",e);}}}})();</script>"
189    )
190}
191
192/// Minimal JSON string literal escaper for embedding a URL path in a
193/// JavaScript string. Handles the only chars the URL grammar permits to
194/// also be JSON-significant: `"` and `\`.
195fn json_string_literal(s: &str) -> String {
196    let mut out = String::with_capacity(s.len() + 2);
197    out.push('"');
198    for ch in s.chars() {
199        match ch {
200            '"' => out.push_str("\\\""),
201            '\\' => out.push_str("\\\\"),
202            c if (c as u32) < 0x20 => {
203                use std::fmt::Write;
204                let _ = write!(&mut out, "\\u{:04x}", c as u32);
205            }
206            c => out.push(c),
207        }
208    }
209    out.push('"');
210    out
211}
212
213#[cfg(test)]
214#[allow(clippy::disallowed_methods)]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn live_reload_endpoint_round_trips() {
220        let lr = LiveReload::new("/__webui/livereload");
221        assert_eq!(lr.endpoint(), "/__webui/livereload");
222    }
223
224    #[test]
225    fn client_script_includes_endpoint() {
226        let lr = LiveReload::new("/__webui/livereload");
227        assert!(lr.client_script().contains("/__webui/livereload"));
228        assert!(lr.client_script().contains("EventSource"));
229        assert!(lr.client_script().contains("reload"));
230    }
231
232    #[test]
233    fn inject_places_script_before_close_body() {
234        let lr = LiveReload::new("/__webui/lr");
235        let html = "<html><body>x</body></html>";
236        let injected = lr.inject(html);
237        let script_idx = injected.find("EventSource").unwrap();
238        let close_idx = injected.find("</body>").unwrap();
239        assert!(script_idx < close_idx);
240    }
241
242    #[test]
243    fn broadcast_reload_does_not_panic_with_no_subscribers() {
244        let lr = LiveReload::new("/x");
245        lr.broadcast_reload();
246        lr.broadcast_error("oops");
247    }
248
249    #[test]
250    fn broadcast_reload_reaches_subscriber() {
251        let lr = LiveReload::new("/x");
252        let mut rx = lr.sender().subscribe();
253        lr.broadcast_reload();
254        let evt = rx.try_recv().unwrap();
255        assert!(matches!(evt, ReloadEvent::Reload));
256    }
257
258    #[test]
259    fn broadcast_error_carries_message() {
260        let lr = LiveReload::new("/x");
261        let mut rx = lr.sender().subscribe();
262        lr.broadcast_error("build failed");
263        let evt = rx.try_recv().unwrap();
264        match evt {
265            ReloadEvent::Error(msg) => assert_eq!(msg, "build failed"),
266            _ => panic!("expected Error event"),
267        }
268    }
269
270    #[test]
271    fn json_string_literal_escapes_quote_and_backslash() {
272        assert_eq!(json_string_literal("/foo"), "\"/foo\"");
273        assert_eq!(json_string_literal("a\"b"), "\"a\\\"b\"");
274        assert_eq!(json_string_literal("a\\b"), "\"a\\\\b\"");
275    }
276}