pax_compiler/design_server/
mod.rs

1use actix::Addr;
2use actix_web::middleware::Logger;
3
4use actix_web::web::Data;
5use actix_web::{get, web, App, HttpRequest, HttpServer, Responder};
6use actix_web::{post, HttpResponse, Result};
7use actix_web_actors::ws;
8use colored::Colorize;
9use pax_generation::{AIModel, PaxAppGenerator};
10use serde_json::json;
11use std::net::TcpListener;
12use std::{env, fs};
13
14use env_logger;
15use std::io::Write;
16
17use crate::helpers::PAX_BADGE;
18use crate::{RunContext, RunTarget};
19use notify::{Error, Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
20use pax_manifest::PaxManifest;
21
22use std::path::{Path, PathBuf};
23use std::str::FromStr;
24use std::sync::{Arc, Mutex};
25use std::time::{SystemTime, UNIX_EPOCH};
26
27use websocket::PrivilegedAgentWebSocket;
28use websocket::SocketMessageAccumulator;
29
30#[allow(unused)]
31mod llm;
32pub mod static_server;
33pub mod websocket;
34
35pub struct AppState {
36    serve_dir: Mutex<PathBuf>,
37    userland_project_root: Mutex<PathBuf>,
38    active_websocket_client: Mutex<Option<Addr<PrivilegedAgentWebSocket>>>,
39    request_id_counter: Mutex<usize>,
40    manifest: Mutex<Option<PaxManifest>>,
41    last_written_timestamp: Mutex<SystemTime>,
42}
43
44impl AppState {
45    pub fn new_empty() -> Self {
46        Self {
47            serve_dir: Mutex::new(PathBuf::new()),
48            userland_project_root: Mutex::new(PathBuf::new()),
49            active_websocket_client: Mutex::new(None),
50            request_id_counter: Mutex::new(0),
51            manifest: Mutex::new(None),
52            last_written_timestamp: Mutex::new(UNIX_EPOCH),
53        }
54    }
55    pub fn new(serve_dir: PathBuf, project_root: PathBuf, manifest: PaxManifest) -> Self {
56        AppState {
57            serve_dir: Mutex::new(serve_dir),
58            userland_project_root: Mutex::new(project_root),
59            active_websocket_client: Mutex::new(None),
60            request_id_counter: Mutex::new(0),
61            manifest: Mutex::new(Some(manifest)),
62            last_written_timestamp: Mutex::new(SystemTime::now()),
63        }
64    }
65
66    fn generate_request_id(&self) -> usize {
67        let mut counter = self.request_id_counter.lock().unwrap();
68        *counter += 1;
69        *counter
70    }
71
72    pub fn update_last_written_timestamp(&self) {
73        let mut last_written = self.last_written_timestamp.lock().unwrap();
74        *last_written = SystemTime::now();
75    }
76}
77
78#[get("/ws")]
79pub async fn web_socket(
80    req: HttpRequest,
81    stream: web::Payload,
82    state: web::Data<AppState>,
83) -> impl Responder {
84    ws::WsResponseBuilder::new(PrivilegedAgentWebSocket::new(state), &req, stream)
85        .frame_size(2_000_000)
86        .start()
87}
88
89#[allow(unused_assignments)]
90pub fn start_server(
91    static_file_path: &str,
92    src_folder_to_watch: &str,
93    manifest: PaxManifest,
94) -> std::io::Result<()> {
95    // Initialize logging
96    std::env::set_var("RUST_LOG", "actix_web=info");
97    env_logger::Builder::from_env(env_logger::Env::default())
98        .format(|buf, record| writeln!(buf, "{} 🍱 Served {}", *PAX_BADGE, record.args()))
99        .init();
100
101    let initial_state = AppState::new(
102        PathBuf::from(static_file_path),
103        PathBuf::from_str(src_folder_to_watch).unwrap(),
104        manifest,
105    );
106    let fs_path = initial_state.serve_dir.lock().unwrap().clone();
107    let state = Data::new(initial_state);
108    let _watcher = setup_file_watcher(state.clone(), src_folder_to_watch)
109        .expect("Failed to setup file watcher");
110
111    // Create a Runtime
112    let runtime = actix_web::rt::System::new().block_on(async {
113        let mut port = 8080;
114        let server = loop {
115            // Check if the port is available
116            if TcpListener::bind(("127.0.0.1", port)).is_ok() {
117                // Log the server details
118                println!(
119                    "{} 🗂️  Serving static files from {}",
120                    *PAX_BADGE,
121                    &fs_path.to_str().unwrap()
122                );
123                let address_msg = format!("http://127.0.0.1:{}", port).blue();
124                let server_running_at_msg = format!("Server running at {}", address_msg).bold();
125                println!("{} 📠 {}", *PAX_BADGE, server_running_at_msg);
126                break HttpServer::new(move || {
127                    App::new()
128                        .wrap(Logger::new("| %s | %U"))
129                        .app_data(state.clone())
130                        .service(ai_page)
131                        .service(ai_submit)
132                        .service(web_socket)
133                        .service(
134                            actix_files::Files::new("/*", fs_path.clone()).index_file("index.html"),
135                        )
136                })
137                .bind(("127.0.0.1", port))
138                .expect("Error binding to address")
139                .workers(2);
140            } else {
141                port += 1; // Try the next port
142            }
143        };
144
145        server.run().await
146    });
147
148    runtime
149}
150
151#[derive(Default)]
152pub enum FileContent {
153    Pax(String),
154    Rust(String),
155    #[default]
156    Unknown,
157}
158
159#[derive(Default)]
160struct WatcherFileChanged {
161    pub contents: FileContent,
162    pub path: String,
163}
164
165impl actix::Message for WatcherFileChanged {
166    type Result = ();
167}
168
169pub fn setup_file_watcher(state: Data<AppState>, path: &str) -> Result<RecommendedWatcher, Error> {
170    let mut watcher = RecommendedWatcher::new(
171        move |res: Result<Event, Error>| match res {
172            Ok(e) => {
173                if let Some(addr) = &*state.active_websocket_client.lock().unwrap() {
174                    let now = SystemTime::now();
175                    // check last written time so we don't spam file changes when we serialize
176                    let last_written = *state.last_written_timestamp.lock().unwrap();
177                    if now
178                        .duration_since(last_written)
179                        .unwrap_or_default()
180                        .as_millis()
181                        > 1000
182                    {
183                        if let EventKind::Modify(_) = e.kind {
184                            if let Some(path) = e.paths.first() {
185                                match fs::read_to_string(path) {
186                                    Ok(contents) => {
187                                        let extension = path.extension();
188                                        let msg = WatcherFileChanged {
189                                            contents: match extension.and_then(|e| e.to_str()) {
190                                                Some("pax") => FileContent::Pax(contents),
191                                                Some("rs") => FileContent::Rust(contents),
192                                                _ => FileContent::Unknown,
193                                            },
194                                            path: path.to_str().unwrap().to_string(),
195                                        };
196                                        addr.do_send(msg);
197                                        state.update_last_written_timestamp();
198                                    }
199                                    Err(_) => (),
200                                }
201                            }
202                        }
203                    }
204                }
205            }
206            Err(e) => {
207                println!("File system watch error: {:?}", e);
208            }
209        },
210        Default::default(),
211    )?;
212    watcher.watch(Path::new(path), RecursiveMode::Recursive)?;
213    Ok(watcher)
214}
215
216#[get("/ai")]
217async fn ai_page() -> Result<HttpResponse> {
218    let html_content = fs::read_to_string("static/ai_chat.html")?;
219    Ok(HttpResponse::Ok()
220        .content_type("text/html")
221        .body(html_content))
222}
223
224#[post("/ai")]
225async fn ai_submit(message: web::Json<AiMessage>, state: web::Data<AppState>) -> HttpResponse {
226    let userland_project_root = state.userland_project_root.lock().unwrap().clone();
227    let claude_api_key = match env::var("ANTHROPIC_API_KEY") {
228        Ok(key) => key,
229        Err(_) => {
230            return HttpResponse::InternalServerError().json(json!({
231                "status": "error",
232                "message": "ANTHROPIC_API_KEY not set in environment"
233            }))
234        }
235    };
236
237    let pax_app_generator = PaxAppGenerator::new(claude_api_key, AIModel::Claude3);
238    let output = userland_project_root.clone().join("src");
239
240    match pax_app_generator
241        .generate_app(&message.message, Some(&output), true)
242        .await
243    {
244        Ok(_) => {
245            match perform_build_and_update_state(&state, userland_project_root.to_str().unwrap()) {
246                Ok(_) => HttpResponse::Ok().json(json!({
247                    "status": "success",
248                    "response": "App generated and built successfully.",
249                })),
250                Err(e) => {
251                    println!("Error performing build and updating state: {:?}", e);
252                    HttpResponse::InternalServerError().json(json!({
253                        "status": "error",
254                        "message": "Failed to build the generated app"
255                    }))
256                }
257            }
258        }
259        Err(e) => {
260            println!("Error generating app: {:?}", e);
261            HttpResponse::InternalServerError().json(json!({
262                "status": "error",
263                "message": "Failed to generate app"
264            }))
265        }
266    }
267}
268
269#[derive(Deserialize)]
270struct AiMessage {
271    message: String,
272}
273
274fn create_designer_run_context() -> RunContext {
275    RunContext {
276        target: RunTarget::Web,
277        project_path: PathBuf::from("../pax-designer".to_string()),
278        verbose: false,
279        should_also_run: false,
280        is_libdev_mode: true,
281        should_run_designer: true,
282        process_child_ids: Arc::new(Mutex::new(vec![])),
283        is_release: false,
284    }
285}
286
287fn perform_build() -> std::io::Result<(PaxManifest, Option<PathBuf>)> {
288    let ctx = create_designer_run_context();
289    crate::perform_build(&ctx).map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
290}
291
292fn perform_build_and_update_state(state: &AppState, folder_to_watch: &str) -> std::io::Result<()> {
293    let (manifest, fs_path) = perform_build()?;
294
295    // Update the state
296    *state.serve_dir.lock().unwrap() = fs_path.expect("serve directory should exist");
297    *state.userland_project_root.lock().unwrap() = PathBuf::from_str(folder_to_watch).unwrap();
298    *state.manifest.lock().unwrap() = Some(manifest);
299
300    Ok(())
301}