orbiton/
dev_server.rs

1// Development server for the Orbit UI framework
2
3use anyhow::Result;
4use futures_util::{future, SinkExt, StreamExt};
5use log::{debug, error, info};
6use std::{
7    net::{IpAddr, Ipv4Addr, SocketAddr},
8    path::{Path, PathBuf},
9    sync::Arc,
10    thread,
11};
12use tokio::net::{TcpListener, TcpStream};
13use tokio::sync::broadcast;
14use tokio_tungstenite::{accept_async, tungstenite::protocol::Message};
15
16use crate::hmr::HmrContext;
17use crate::hmr_inject::{get_hmr_client_js, is_html_file, process_html_file};
18
19/// Development server
20pub struct DevServer {
21    /// Port to use for the server
22    port: u16,
23    /// Project directory
24    project_dir: PathBuf,
25    /// Server thread handle
26    #[allow(dead_code)]
27    thread_handle: Option<thread::JoinHandle<()>>,
28    /// Broadcast channel for sending updates to connected clients
29    tx: Option<broadcast::Sender<String>>,
30    /// Use beta toolchain for building and testing
31    use_beta: bool,
32    /// HMR context for tracking changed modules
33    hmr_context: Arc<HmrContext>,
34}
35
36impl Clone for DevServer {
37    fn clone(&self) -> Self {
38        Self {
39            port: self.port,
40            project_dir: self.project_dir.clone(),
41            thread_handle: None, // Don't clone the thread handle
42            tx: self.tx.clone(),
43            use_beta: self.use_beta,
44            hmr_context: Arc::clone(&self.hmr_context),
45        }
46    }
47}
48
49impl DevServer {
50    /// Create a new development server
51    #[allow(dead_code)] // Used in tests and maintenance operations
52    pub fn new(port: u16, project_dir: &Path) -> Result<Self> {
53        let (tx, _) = broadcast::channel(16);
54        let hmr_context = Arc::new(HmrContext::new(project_dir.to_owned()));
55
56        Ok(Self {
57            port,
58            project_dir: project_dir.to_owned(),
59            thread_handle: None,
60            tx: Some(tx),
61            use_beta: false,
62            hmr_context,
63        })
64    }
65
66    /// Create a new development server with optional beta toolchain support
67    pub fn new_with_options(port: u16, project_dir: &Path, use_beta: bool) -> Result<Self> {
68        let (tx, _) = broadcast::channel(16);
69        let hmr_context = Arc::new(HmrContext::new(project_dir.to_owned()));
70
71        Ok(Self {
72            port,
73            project_dir: project_dir.to_owned(),
74            thread_handle: None,
75            tx: Some(tx),
76            use_beta,
77            hmr_context,
78        })
79    }
80
81    /// Check if the dev server is using beta toolchain
82    pub fn is_using_beta(&self) -> bool {
83        self.use_beta
84    }
85
86    /// Get the server port
87    #[allow(dead_code)] // Used in tests and maintenance operations
88    pub fn port(&self) -> u16 {
89        self.port
90    }
91
92    /// Get the HMR context
93    pub fn hmr_context(&self) -> &Arc<HmrContext> {
94        &self.hmr_context
95    }
96
97    /// Start the development server
98    pub fn start(&mut self) -> Result<&thread::JoinHandle<()>> {
99        let port = self.port;
100        let project_dir = self.project_dir.clone();
101        let tx = self.tx.take().expect("Missing broadcast channel");
102
103        let handle = thread::spawn(move || {
104            // Set up the Tokio runtime
105            let rt = tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime");
106
107            rt.block_on(async {
108                // Start WebSocket server
109                let ws_rx = tx.subscribe();
110                let ws_handle = tokio::spawn(Self::run_websocket_server(port, ws_rx));
111
112                // Start HTTP server
113                let server = tiny_http::Server::http(format!("0.0.0.0:{port}"))
114                    .expect("Failed to start HTTP server");
115
116                info!("Development server started on port {port}");
117                info!("WebSocket server started on port {}", port + 1);
118
119                let _broadcast_tx = tx; // Keep tx alive
120
121                for request in server.incoming_requests() {
122                    debug!("Received request: {url:?}", url = request.url());
123
124                    // Special handling for HMR client script
125                    if request.url() == "/__orbit_hmr_client.js" {
126                        debug!("Serving HMR client script");
127                        let response = tiny_http::Response::from_string(get_hmr_client_js())
128                            .with_header(
129                                tiny_http::Header::from_bytes(
130                                    &b"Content-Type"[..],
131                                    &b"application/javascript"[..],
132                                )
133                                .unwrap(),
134                            );
135                        let _ = request.respond(response);
136                        continue;
137                    }
138
139                    // Handle static files
140                    let url = request.url().trim_start_matches('/');
141                    let file_path = if url.is_empty() {
142                        project_dir.join("index.html")
143                    } else {
144                        project_dir.join(url)
145                    };
146
147                    if file_path.exists() && file_path.is_file() {
148                        // Special handling for HTML files to inject HMR client
149                        if is_html_file(&file_path) {
150                            debug!("Processing HTML file: {:?}", file_path);
151                            match process_html_file(&file_path, port) {
152                                Ok(content) => {
153                                    let response = tiny_http::Response::from_data(content)
154                                        .with_header(
155                                            tiny_http::Header::from_bytes(
156                                                &b"Content-Type"[..],
157                                                &b"text/html"[..],
158                                            )
159                                            .unwrap(),
160                                        );
161                                    let _ = request.respond(response);
162                                }
163                                Err(e) => {
164                                    error!("Failed to process HTML file: {}", e);
165                                    // Fall back to serving the file without injection
166                                    let file = std::fs::File::open(&file_path)
167                                        .expect("Failed to open file");
168                                    let response = tiny_http::Response::from_file(file);
169                                    let _ = request.respond(response);
170                                }
171                            }
172                        } else {
173                            // Serve non-HTML files normally
174                            let file =
175                                std::fs::File::open(&file_path).expect("Failed to open file");
176                            let response = tiny_http::Response::from_file(file);
177                            let _ = request.respond(response);
178                        }
179                    } else {
180                        // File not found, return 404
181                        let response = tiny_http::Response::from_string("File not found")
182                            .with_status_code(404);
183                        let _ = request.respond(response);
184                    }
185                }
186
187                // Wait for WebSocket server to finish
188                let _ = ws_handle.await;
189            });
190        });
191
192        self.thread_handle = Some(handle);
193        Ok(self.thread_handle.as_ref().unwrap())
194    }
195
196    /// Send an update to all connected WebSocket clients
197    pub fn broadcast_update(&self, message: String) -> Result<()> {
198        if let Some(tx) = &self.tx {
199            tx.send(message)
200                .map_err(|e| anyhow::anyhow!("Failed to broadcast message: {}", e))?;
201        }
202        Ok(())
203    }
204
205    /// Trigger an HMR update for specific modules
206    pub fn send_hmr_update(&self, modules: Vec<String>) -> Result<()> {
207        let message = serde_json::json!({
208            "type": "hmr",
209            "modules": modules
210        })
211        .to_string();
212
213        self.broadcast_update(message)
214    }
215
216    /// Trigger a full page reload for all clients
217    pub fn send_reload_command(&self) -> Result<()> {
218        let message = serde_json::json!({
219            "type": "reload"
220        })
221        .to_string();
222
223        self.broadcast_update(message)
224    }
225
226    /// Send rebuild status to all clients
227    pub fn send_rebuild_status(&self, status: &str) -> Result<()> {
228        let message = serde_json::json!({
229            "type": "rebuild",
230            "status": status
231        })
232        .to_string();
233
234        self.broadcast_update(message)
235    }
236
237    async fn handle_websocket_connection(
238        ws_stream: tokio_tungstenite::WebSocketStream<TcpStream>,
239        addr: SocketAddr,
240        mut rx: broadcast::Receiver<String>,
241    ) {
242        info!("WebSocket connection established: {addr}");
243        let (mut ws_sender, mut ws_receiver) = ws_stream.split();
244
245        // Send initial connection acknowledgment
246        let hello_msg = serde_json::json!({
247            "type": "hello",
248            "message": "Orbit HMR connected"
249        })
250        .to_string();
251
252        if let Err(e) = ws_sender.send(Message::Text(hello_msg)).await {
253            error!("Error sending hello message: {e}");
254            return;
255        }
256
257        let send_task = tokio::spawn(async move {
258            while let Ok(msg) = rx.recv().await {
259                ws_sender
260                    .send(Message::Text(msg))
261                    .await
262                    .unwrap_or_else(|e| error!("Error sending message: {e}"));
263            }
264        });
265
266        let recv_task = tokio::spawn(async move {
267            while let Some(msg) = ws_receiver.next().await {
268                if let Ok(msg) = msg {
269                    if msg.is_close() {
270                        break;
271                    }
272
273                    // Handle incoming messages from client
274                    if let Message::Text(text) = msg {
275                        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&text) {
276                            if let Some(msg_type) = json.get("type").and_then(|t| t.as_str()) {
277                                match msg_type {
278                                    "register" => {
279                                        if let Some(path) = json.get("url").and_then(|p| p.as_str())
280                                        {
281                                            debug!("Client registered for path: {}", path);
282                                            // Could store client info in a map for targeted updates
283                                        }
284                                    }
285                                    "hmr_ready" => {
286                                        debug!("Client reported HMR ready state");
287                                    }
288                                    _ => debug!("Received unknown message type: {}", msg_type),
289                                }
290                            }
291                        }
292                    }
293                }
294            }
295        });
296
297        future::select(send_task, recv_task).await;
298        info!("WebSocket connection closed: {addr}");
299    }
300
301    /// Start the WebSocket server
302    async fn run_websocket_server(port: u16, rx: broadcast::Receiver<String>) -> Result<()> {
303        let addr = (IpAddr::V4(Ipv4Addr::LOCALHOST), port + 1);
304        let listener = TcpListener::bind(addr).await?;
305        info!("WebSocket server listening on: localhost:{}", port + 1);
306
307        while let Ok((stream, addr)) = listener.accept().await {
308            let ws_stream = accept_async(stream).await?;
309            let rx = rx.resubscribe();
310
311            tokio::spawn(async move {
312                Self::handle_websocket_connection(ws_stream, addr, rx).await;
313            });
314        }
315        Ok(())
316    }
317}
orbiton/dev_server.rs

orbiton/
dev_server.rs
