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