pax_designtime/
privileged_agent.rs

1use crate::messages::{
2    AgentMessage, ComponentSerializationRequest, DevClientResponse, LLMRequest,
3    LoadFileToStaticDirRequest, UserlandSourceUpdateRequest,
4};
5use anyhow::{anyhow, Result};
6use ewebsock::{WsEvent, WsMessage};
7use pax_manifest::ComponentDefinition;
8use std::time::Duration;
9#[cfg(not(target_arch = "wasm32"))]
10use std::time::Instant;
11use url::Url;
12#[cfg(target_arch = "wasm32")]
13use web_time::Instant;
14
15const WEBSOCKET_RECONNECT_INITIAL_DELAY: Duration = Duration::from_millis(500);
16const WEBSOCKET_RECONNECT_MAX_DELAY: Duration = Duration::from_secs(5);
17
18#[derive(Debug)]
19pub(crate) enum ConnectionEvent {
20    Opened,
21    Message(AgentMessage),
22    Disconnected,
23}
24
25pub struct WebSocketConnection {
26    url: String,
27    sender: Option<ewebsock::WsSender>,
28    recver: Option<ewebsock::WsReceiver>,
29    label: String,
30    pub alive: bool,
31    allow_reconnect: bool,
32    connecting: bool,
33    next_reconnect_at: Option<Instant>,
34    reconnect_delay: Duration,
35}
36
37impl WebSocketConnection {
38    pub fn new(
39        addr: &str,
40        versioning_prefix: Option<&str>,
41        label: impl Into<String>,
42    ) -> Result<Self> {
43        let url = build_socket_url(addr, versioning_prefix)?;
44        let (sender, recver) = connect_socket(&url)?;
45
46        Ok(Self {
47            url,
48            sender: Some(sender),
49            recver: Some(recver),
50            label: label.into(),
51            // The browser WebSocket API rejects sends while CONNECTING. Revision messages are
52            // retained by RevisionGate and replayed after ConnectionEvent::Opened.
53            alive: false,
54            allow_reconnect: true,
55            connecting: true,
56            next_reconnect_at: None,
57            reconnect_delay: WEBSOCKET_RECONNECT_INITIAL_DELAY,
58        })
59    }
60
61    /// Creates an explicitly disconnected transport. Native debug cartridges
62    /// use this when no design-server address was supplied instead of probing a
63    /// magical localhost endpoint forever.
64    pub fn offline(label: impl Into<String>) -> Self {
65        Self {
66            url: String::new(),
67            sender: None,
68            recver: None,
69            label: label.into(),
70            alive: false,
71            allow_reconnect: false,
72            connecting: false,
73            next_reconnect_at: None,
74            reconnect_delay: WEBSOCKET_RECONNECT_INITIAL_DELAY,
75        }
76    }
77
78    pub fn send_agent_message(&mut self, message: &AgentMessage) -> Result<()> {
79        if !self.alive || self.connecting {
80            return Err(anyhow!("design-server socket is not open"));
81        }
82        let msg_bytes = rmp_serde::to_vec(message)?;
83        self.sender()
84            .ok_or_else(|| anyhow!("design-server socket is not connected"))?
85            .send(ewebsock::WsMessage::Binary(msg_bytes));
86        Ok(())
87    }
88
89    pub fn send_component_update(&mut self, component: &ComponentDefinition) -> Result<()> {
90        let component_bytes = rmp_serde::to_vec(component)?;
91        self.send_connected(
92            AgentMessage::ComponentSerializationRequest(ComponentSerializationRequest {
93                component_bytes,
94            }),
95            "couldn't send component update: connection to design-server was lost",
96        )
97    }
98
99    pub fn send_llm_request(&mut self, llm_request: LLMRequest) -> Result<()> {
100        self.send_connected(
101            AgentMessage::LLMRequest(llm_request),
102            "couldn't send LLM request: connection to pub pax was lost",
103        )
104    }
105
106    pub fn send_file_to_static_dir(&mut self, name: &str, data: Vec<u8>) -> Result<()> {
107        self.send_connected(
108            AgentMessage::LoadFileToStaticDirRequest(LoadFileToStaticDirRequest {
109                name: name.to_owned(),
110                data,
111            }),
112            "couldn't send file: connection to design-server was lost",
113        )
114    }
115
116    pub fn send_dev_client_response(&mut self, response: DevClientResponse) -> Result<()> {
117        self.send_connected(
118            AgentMessage::DevClientResponse(response),
119            "couldn't send dev response: connection to design-server was lost",
120        )
121    }
122
123    pub fn send_userland_source_update_request(
124        &mut self,
125        request: UserlandSourceUpdateRequest,
126    ) -> Result<()> {
127        self.send_connected(
128            AgentMessage::UserlandSourceUpdateRequest(request),
129            "couldn't send userland source update: connection to design-server was lost",
130        )
131    }
132
133    fn send_connected(&mut self, message: AgentMessage, disconnected_error: &str) -> Result<()> {
134        if !self.alive {
135            return Err(anyhow!(disconnected_error.to_owned()));
136        }
137        self.send_agent_message(&message)
138    }
139
140    pub(crate) fn handle_recv(&mut self) -> Result<Vec<ConnectionEvent>> {
141        self.reconnect_if_needed();
142
143        let mut connection_events = vec![];
144        while let Some(event) = self.recver.as_ref().and_then(|recver| recver.try_recv()) {
145            match event {
146                WsEvent::Opened => {
147                    self.alive = true;
148                    self.connecting = false;
149                    self.next_reconnect_at = None;
150                    self.reconnect_delay = WEBSOCKET_RECONNECT_INITIAL_DELAY;
151                    connection_events.push(ConnectionEvent::Opened);
152                }
153                WsEvent::Message(message) => match message {
154                    WsMessage::Binary(msg_bytes) => {
155                        let msg: AgentMessage = match rmp_serde::from_slice(&msg_bytes) {
156                            Ok(msg) => msg,
157                            Err(err) => {
158                                log::warn!(
159                                    "{} received invalid websocket message: {err}",
160                                    self.label
161                                );
162                                continue;
163                            }
164                        };
165                        if let AgentMessage::DisconnectNotification(notification) = &msg {
166                            self.allow_reconnect = notification.allow_reconnect;
167                            if !notification.allow_reconnect {
168                                log::info!(
169                                    "{} reconnect disabled by server: {}",
170                                    self.label,
171                                    notification.reason
172                                );
173                            }
174                        }
175                        connection_events.push(ConnectionEvent::Message(msg));
176                    }
177                    WsMessage::Ping(data) => {
178                        if let Some(sender) = self.sender() {
179                            sender.send(WsMessage::Pong(data));
180                        }
181                    }
182                    WsMessage::Pong(_) | WsMessage::Text(_) | WsMessage::Unknown(_) => {}
183                },
184                WsEvent::Error(e) => {
185                    log::warn!("{} web socket error: {e}", self.label);
186                    self.schedule_reconnect();
187                    connection_events.push(ConnectionEvent::Disconnected);
188                }
189                WsEvent::Closed => {
190                    log::warn!("{} web socket was closed", self.label);
191                    self.schedule_reconnect();
192                    connection_events.push(ConnectionEvent::Disconnected);
193                }
194            }
195        }
196        Ok(connection_events)
197    }
198
199    fn sender(&mut self) -> Option<&mut ewebsock::WsSender> {
200        self.sender.as_mut()
201    }
202
203    fn schedule_reconnect(&mut self) {
204        self.sender = None;
205        self.recver = None;
206        self.alive = false;
207        self.connecting = false;
208        if !self.allow_reconnect {
209            self.next_reconnect_at = None;
210            return;
211        }
212        if self.next_reconnect_at.is_none() {
213            self.next_reconnect_at = Some(Instant::now() + self.reconnect_delay);
214            self.reconnect_delay = self
215                .reconnect_delay
216                .saturating_mul(2)
217                .min(WEBSOCKET_RECONNECT_MAX_DELAY);
218        }
219    }
220
221    fn reconnect_if_needed(&mut self) {
222        if !self.allow_reconnect {
223            return;
224        }
225        if self.alive || self.connecting {
226            return;
227        }
228        let Some(next_reconnect_at) = self.next_reconnect_at else {
229            return;
230        };
231        if Instant::now() < next_reconnect_at {
232            return;
233        }
234
235        self.next_reconnect_at = None;
236        match connect_socket(&self.url) {
237            Ok((sender, recver)) => {
238                self.sender = Some(sender);
239                self.recver = Some(recver);
240                self.alive = false;
241                self.connecting = true;
242                self.next_reconnect_at = None;
243                log::info!("{} reconnecting to {}", self.label, self.url);
244            }
245            Err(err) => {
246                log::warn!("{} reconnect failed: {err}", self.label);
247                self.schedule_reconnect();
248            }
249        }
250    }
251
252    pub fn shutdown_permanently(&mut self) {
253        self.allow_reconnect = false;
254        if let Some(sender) = self.sender.as_mut() {
255            let _ = sender.close();
256        }
257        self.sender = None;
258        self.recver = None;
259        self.alive = false;
260        self.connecting = false;
261        self.next_reconnect_at = None;
262    }
263}
264
265fn build_socket_url(addr: &str, versioning_prefix: Option<&str>) -> Result<String> {
266    let mut url = Url::parse(addr).map_err(|e| anyhow!("Invalid URL: {}", e))?;
267
268    let ws_scheme = match url.scheme() {
269        "http" => "ws",
270        "https" => "wss",
271        "ws" | "wss" => url.scheme(),
272        _ => return Err(anyhow!("Unsupported URL scheme: {}", url.scheme())),
273    }
274    .to_owned();
275
276    url.set_scheme(&ws_scheme)
277        .map_err(|_| anyhow!("Failed to set URL scheme"))?;
278
279    let versioning_prefix = versioning_prefix.unwrap_or("");
280    let new_path = format!("{}/ws", versioning_prefix);
281    url.set_path(&new_path);
282
283    Ok(url.to_string())
284}
285
286fn connect_socket(url: &str) -> Result<(ewebsock::WsSender, ewebsock::WsReceiver)> {
287    ewebsock::connect_with_wakeup(
288        url.to_owned(),
289        ewebsock::Options::default(),
290        wake_designtime_loop,
291    )
292    .map_err(|err| anyhow!("Couldn't create socket connection: {err}"))
293}
294
295#[cfg(target_arch = "wasm32")]
296fn wake_designtime_loop() {
297    if let Some(window) = web_sys::window() {
298        if let Ok(event) = web_sys::Event::new("pax-designtime-wakeup") {
299            let _ = window.dispatch_event(&event);
300        }
301    }
302}
303
304#[cfg(not(target_arch = "wasm32"))]
305fn wake_designtime_loop() {}
306
307#[cfg(test)]
308mod tests {
309    use super::{
310        build_socket_url, ConnectionEvent, WebSocketConnection, WEBSOCKET_RECONNECT_INITIAL_DELAY,
311        WEBSOCKET_RECONNECT_MAX_DELAY,
312    };
313    use crate::messages::{AgentMessage, FileChangedNotification};
314    use ewebsock::{WsEvent, WsMessage};
315
316    fn test_connection(recver: ewebsock::WsReceiver) -> WebSocketConnection {
317        WebSocketConnection {
318            url: "ws://127.0.0.1:1/ws".to_string(),
319            sender: None,
320            recver: Some(recver),
321            label: "test".to_string(),
322            alive: true,
323            allow_reconnect: true,
324            connecting: false,
325            next_reconnect_at: None,
326            reconnect_delay: WEBSOCKET_RECONNECT_INITIAL_DELAY,
327        }
328    }
329
330    #[test]
331    fn builds_websocket_url_from_http_origin() {
332        assert_eq!(
333            build_socket_url("http://127.0.0.1:8080", None).unwrap(),
334            "ws://127.0.0.1:8080/ws"
335        );
336    }
337
338    #[test]
339    fn preserves_version_prefix_when_building_websocket_url() {
340        assert_eq!(
341            build_socket_url("https://pub.pax.dev", Some("/v0")).unwrap(),
342            "wss://pub.pax.dev/v0/ws"
343        );
344    }
345
346    #[test]
347    fn ignores_invalid_websocket_messages() {
348        let (recver, on_event) = ewebsock::WsReceiver::new();
349        let _ = on_event(WsEvent::Message(WsMessage::Binary(vec![0xc1])));
350        let mut connection = test_connection(recver);
351        let messages = connection.handle_recv().unwrap();
352
353        assert!(messages.is_empty());
354        assert!(connection.alive);
355    }
356
357    #[test]
358    fn forwards_decoded_messages_without_applying_policy() {
359        let (recver, on_event) = ewebsock::WsReceiver::new();
360        let message = AgentMessage::ProjectFileChangedNotification(FileChangedNotification {});
361        let _ = on_event(WsEvent::Message(WsMessage::Binary(
362            rmp_serde::to_vec(&message).unwrap(),
363        )));
364        let mut connection = test_connection(recver);
365
366        let messages = connection.handle_recv().unwrap();
367
368        assert!(matches!(
369            messages.as_slice(),
370            [ConnectionEvent::Message(
371                AgentMessage::ProjectFileChangedNotification(_)
372            )]
373        ));
374    }
375
376    #[test]
377    fn schedules_reconnect_after_close_without_error() {
378        let (recver, on_event) = ewebsock::WsReceiver::new();
379        let _ = on_event(WsEvent::Closed);
380        let mut connection = test_connection(recver);
381
382        let messages = connection.handle_recv().unwrap();
383
384        assert!(matches!(
385            messages.as_slice(),
386            [ConnectionEvent::Disconnected]
387        ));
388        assert!(!connection.alive);
389        assert!(!connection.connecting);
390        assert!(connection.next_reconnect_at.is_some());
391    }
392
393    #[test]
394    fn schedules_reconnect_after_error_without_error() {
395        let (recver, on_event) = ewebsock::WsReceiver::new();
396        let _ = on_event(WsEvent::Error(String::new()));
397        let mut connection = test_connection(recver);
398
399        let messages = connection.handle_recv().unwrap();
400
401        assert!(matches!(
402            messages.as_slice(),
403            [ConnectionEvent::Disconnected]
404        ));
405        assert!(!connection.alive);
406        assert!(!connection.connecting);
407        assert!(connection.next_reconnect_at.is_some());
408    }
409
410    #[test]
411    fn reconnect_attempts_use_capped_backoff() {
412        let (recver, _on_event) = ewebsock::WsReceiver::new();
413        let mut connection = test_connection(recver);
414
415        connection.schedule_reconnect();
416        assert_eq!(
417            connection.reconnect_delay,
418            WEBSOCKET_RECONNECT_INITIAL_DELAY.saturating_mul(2)
419        );
420
421        for _ in 0..8 {
422            connection.next_reconnect_at = None;
423            connection.schedule_reconnect();
424        }
425        assert_eq!(connection.reconnect_delay, WEBSOCKET_RECONNECT_MAX_DELAY);
426    }
427
428    #[test]
429    fn offline_transport_never_reconnects() {
430        let mut connection = WebSocketConnection::offline("test-offline");
431
432        assert!(connection.handle_recv().unwrap().is_empty());
433        assert!(!connection.alive);
434        assert!(!connection.allow_reconnect);
435        assert!(connection.next_reconnect_at.is_none());
436    }
437
438    #[test]
439    fn connecting_transport_rejects_sends_until_opened() {
440        let (recver, on_event) = ewebsock::WsReceiver::new();
441        let mut connection = test_connection(recver);
442        connection.alive = false;
443        connection.connecting = true;
444
445        let err = connection
446            .send_agent_message(&AgentMessage::ProjectFileChangedNotification(
447                FileChangedNotification {},
448            ))
449            .unwrap_err();
450        assert_eq!(err.to_string(), "design-server socket is not open");
451
452        let _ = on_event(WsEvent::Opened);
453        let events = connection.handle_recv().unwrap();
454        assert!(matches!(events.as_slice(), [ConnectionEvent::Opened]));
455        assert!(connection.alive);
456        assert!(!connection.connecting);
457    }
458}
pax_designtime/privileged_agent.rs

pax_designtime/
privileged_agent.rs
