pax_designtime/
lib.rs

1use std::any::Any;
2use std::cell::RefCell;
3use std::collections::{HashMap, HashSet};
4use std::net::{Ipv4Addr, SocketAddr, ToSocketAddrs};
5use std::rc::Rc;
6
7pub mod orm;
8pub mod privileged_agent;
9
10pub mod messages;
11pub mod serde_pax;
12
13use messages::LLMRequest;
14use orm::{MessageType, ReloadType};
15use pax_manifest::pax_runtime_api::Property;
16use privileged_agent::WebSocketConnection;
17use url::Url;
18
19use core::fmt::Debug;
20
21pub use pax_manifest;
22use pax_manifest::{
23    server::*, ComponentDefinition, PaxManifest, TypeId, UniqueTemplateNodeIdentifier,
24};
25use reqwasm::http::Response;
26pub use serde_pax::error::{Error, Result};
27pub use serde_pax::se::{to_pax, Serializer};
28
29pub const INITIAL_MANIFEST_FILE_NAME: &str = "initial-manifest.json";
30
31type Factories = HashMap<String, Box<fn(ComponentDefinition) -> Box<dyn Any>>>;
32use crate::orm::PaxManifestORM;
33
34pub struct DesigntimeManager {
35    orm: PaxManifestORM,
36    factories: Factories,
37    priv_agent_connection: Rc<RefCell<WebSocketConnection>>,
38    pub_pax_connection: Rc<RefCell<WebSocketConnection>>,
39    project_query: Option<String>,
40    response_queue: Rc<RefCell<Vec<DesigntimeResponseMessage>>>,
41    last_rendered_manifest_version: Property<usize>,
42    pub publish_state: Property<Option<PublishResponse>>,
43}
44
45pub enum DesigntimeResponseMessage {
46    LLMResponse(ComponentDefinition),
47    PublishResponse(PublishResponse),
48}
49
50impl Debug for DesigntimeManager {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("DesigntimeManager").finish()
53    }
54}
55
56const VERSION_PREFIX: &str = "/v0";
57const ENDPOINT_PUBLISH: &str = "/v0/publish";
58
59const PROD_PUB_PAX_SERVER: &str = "https://pub.pax.dev";
60
61fn get_server_base_url() -> String {
62    // Fetch the environment variable or use the default value
63    option_env!("PUB_PAX_SERVER")
64        .unwrap_or(PROD_PUB_PAX_SERVER)
65        .to_string()
66}
67
68impl DesigntimeManager {
69    pub fn get_last_rendered_manifest_version(&self) -> Property<usize> {
70        self.last_rendered_manifest_version.clone()
71    }
72
73    pub fn set_last_rendered_manifest_version(&self, version: usize) {
74        self.last_rendered_manifest_version.set(version);
75    }
76
77    pub fn get_llm_messages(&mut self, request_id: u64) -> Vec<String> {
78        self.orm.get_messages(request_id)
79    }
80
81    pub fn new_with_local_addr(manifest: PaxManifest, local_addr: &str) -> Self {
82        let priv_agent = Rc::new(RefCell::new(
83            WebSocketConnection::new(local_addr, None)
84                .expect("couldn't connect to privileged agent"),
85        ));
86
87        let address: String = get_server_base_url();
88        let pub_pax = Rc::new(RefCell::new(
89            WebSocketConnection::new(&address, Some(VERSION_PREFIX))
90                .expect("couldn't connect to privileged agent"),
91        ));
92
93        let orm = PaxManifestORM::new(manifest);
94        let factories = HashMap::new();
95        DesigntimeManager {
96            orm,
97            factories,
98            priv_agent_connection: priv_agent,
99            pub_pax_connection: pub_pax,
100            project_query: None,
101            response_queue: Rc::new(RefCell::new(Vec::new())),
102            last_rendered_manifest_version: Property::new(0),
103            publish_state: Default::default(),
104        }
105    }
106    pub fn new(manifest: PaxManifest) -> Self {
107        Self::new_with_local_addr(manifest, "ws://localhost:8080")
108    }
109
110    pub fn set_project(&mut self, project_query: String) {
111        self.project_query = Some(project_query);
112    }
113
114    pub fn send_file_to_static_dir(&self, name: &str, data: Vec<u8>) -> anyhow::Result<()> {
115        self.priv_agent_connection
116            .borrow_mut()
117            .send_file_to_static_dir(name, data)?;
118        Ok(())
119    }
120
121    pub fn get_new_message_listener(&self) -> Property<Vec<MessageType>> {
122        self.orm.new_message.clone()
123    }
124
125    pub fn get_manifest_loaded_from_server_prop(&self) -> Property<bool> {
126        self.orm.manifest_loaded_from_server.clone()
127    }
128
129    pub fn send_component_update(&mut self, type_id: &TypeId) -> anyhow::Result<()> {
130        self.orm.send_component_update(type_id);
131        let component = self.orm.get_component(type_id)?;
132        self.priv_agent_connection
133            .borrow_mut()
134            .send_component_update(component)?;
135
136        Ok(())
137    }
138
139    pub fn llm_request(&mut self, prompt: &str, request_id: u64) -> anyhow::Result<()> {
140        let manifest = self.orm.get_manifest().clone();
141
142        let llm_request = LLMRequest {
143            manifest: manifest.clone(),
144            prompt: prompt.to_string(),
145            request_id,
146        };
147
148        self.pub_pax_connection
149            .borrow_mut()
150            .send_llm_request(llm_request)?;
151
152        Ok(())
153    }
154
155    pub fn publish_project(&mut self) {
156        let manifest = self.orm.get_manifest().clone();
157
158        let publish_request = PublishRequest {
159            manifest: manifest.clone(),
160        };
161
162        let queue_cloned = self.response_queue.clone();
163
164        wasm_bindgen_futures::spawn_local(async move {
165            let url = get_server_base_url() + ENDPOINT_PUBLISH;
166
167            let response = reqwasm::http::Request::post(&url)
168                .header("Content-Type", "application/json")
169                .body(serde_json::to_string(&publish_request).unwrap())
170                .send()
171                .await;
172
173            log::info!("response_text: {:?}", response);
174
175            let pub_response = match response {
176                Ok(resp) => {
177                    let pub_response: PublishResponse = resp.json().await.unwrap();
178
179                    let pub_response = if let PublishResponse::Success(prs) = pub_response {
180                        prs
181                    } else {
182                        unimplemented!()
183                    };
184                    log::info!("publish success: {:?}", &pub_response.pull_request_url);
185                    PublishResponse::Success(pub_response)
186                }
187                Err(msg) => PublishResponse::Error(ResponseError {
188                    message: msg.to_string(),
189                }),
190            };
191
192            queue_cloned
193                .borrow_mut()
194                .push(DesigntimeResponseMessage::PublishResponse(pub_response));
195        });
196    }
197
198    pub fn add_factory(
199        &mut self,
200        type_id: String,
201        factory: Box<fn(ComponentDefinition) -> Box<dyn Any>>,
202    ) {
203        self.factories.insert(type_id, factory);
204    }
205
206    pub fn get_manifest(&self) -> &PaxManifest {
207        self.orm.get_manifest()
208    }
209
210    pub fn take_reload_queue(&mut self) -> HashSet<ReloadType> {
211        self.orm.take_reload_queue()
212    }
213
214    pub fn reload(&mut self) {
215        self.orm.insert_reload(ReloadType::Tree);
216        self.orm.increment_manifest_version();
217    }
218
219    pub fn set_userland_root_component_type_id(&mut self, type_id: &TypeId) {
220        self.orm.set_userland_root_component_type_id(type_id);
221        self.orm.increment_manifest_version();
222        self.orm.insert_reload(ReloadType::Tree);
223    }
224
225    pub fn get_last_written_manifest_version(&self) -> Property<usize> {
226        self.orm.get_manifest_version()
227    }
228
229    pub fn get_orm(&self) -> &PaxManifestORM {
230        &self.orm
231    }
232
233    pub fn get_orm_mut(&mut self) -> &mut PaxManifestORM {
234        &mut self.orm
235    }
236
237    pub fn handle_recv(&mut self) -> anyhow::Result<()> {
238        self.priv_agent_connection
239            .borrow_mut()
240            .handle_recv(&mut self.orm)?;
241
242        self.pub_pax_connection
243            .borrow_mut()
244            .handle_recv(&mut self.orm)?;
245
246        let response_queue = {
247            let mut queue = self.response_queue.borrow_mut();
248            queue.drain(..).collect::<Vec<DesigntimeResponseMessage>>()
249        };
250        for response in response_queue {
251            self.handle_response(response);
252        }
253        Ok(())
254    }
255
256    pub fn handle_response(&mut self, response: DesigntimeResponseMessage) {
257        match response {
258            DesigntimeResponseMessage::LLMResponse(component) => {
259                log::info!("handling LLM response");
260                let _ = self.orm.swap_main_component(component).map_err(|e| {
261                    log::error!("Error swapping main component for LLM response: {:?}", e);
262                });
263            }
264            DesigntimeResponseMessage::PublishResponse(response) => {
265                log::info!("received publish response");
266                self.publish_state.set(Some(response));
267            }
268        }
269    }
270}
271
272pub enum Args {}
273
274pub struct NodeWithBounds {
275    pub uni: UniqueTemplateNodeIdentifier,
276    pub x: f64,
277    pub y: f64,
278}
pax_designtime/lib.rs

pax_designtime/
lib.rs
