Skip to main content

mcpr_core/proxy/
widgets.rs

1//! Widget HTML bundle serving + asset discovery.
2//!
3//! MCP servers can expose UI resources at `ui://widget/<name>`. This module
4//! serves the bundled HTML either from a local directory (`WidgetSource::Static`)
5//! or by reverse-proxying a dev server (`WidgetSource::Proxy`). Asset URLs in
6//! the served HTML are rewritten to point at the proxy so they resolve
7//! through the tunnel instead of the sandbox origin.
8//!
9//! Used by the pipeline's `WidgetOverlayMiddleware` (for `resources/read`
10//! overlays) and by the widget static routes (`/widgets/<name>.html`,
11//! `/widgets`, and arbitrary asset GETs).
12
13use axum::{
14    http::{HeaderMap, StatusCode, header},
15    response::{IntoResponse, Response},
16};
17use std::path::{Path, PathBuf};
18use std::time::Duration;
19
20use crate::proxy::proxy_state::ProxyState;
21
22// ── Types ───────────────────────────────────────────────
23
24#[derive(Clone)]
25pub enum WidgetSource {
26    /// Reverse proxy to a running server (e.g., http://localhost:4444)
27    Proxy(String),
28    /// Serve static files from a directory (e.g., ./widgets/dist)
29    Static(String),
30}
31
32// ── Asset serving ───────────────────────────────────────
33
34/// Serve a widget asset by path. Called from proxy's catch-all handler for static asset requests.
35pub async fn serve_widget_asset(state: &ProxyState, path: &str) -> Response {
36    match &state.widget_source {
37        Some(WidgetSource::Proxy(base_url)) => {
38            let url = format!("{}{}", base_url.trim_end_matches('/'), path);
39            match state
40                .upstream
41                .http_client
42                .get(&url)
43                .timeout(Duration::from_secs(10))
44                .send()
45                .await
46            {
47                Ok(resp) => {
48                    let status = resp.status().as_u16();
49                    let mut headers = HeaderMap::new();
50                    if let Some(ct) = resp.headers().get(header::CONTENT_TYPE) {
51                        headers.insert(header::CONTENT_TYPE, ct.clone());
52                    }
53                    headers.insert(header::CACHE_CONTROL, "no-cache".parse().unwrap());
54                    let status_code =
55                        StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY);
56                    let bytes = resp.bytes().await.unwrap_or_default();
57                    (status_code, headers, bytes).into_response()
58                }
59                Err(e) => {
60                    (StatusCode::BAD_GATEWAY, format!("Widget proxy error: {e}")).into_response()
61                }
62            }
63        }
64        Some(WidgetSource::Static(dir)) => {
65            let file_path = PathBuf::from(dir).join(path.trim_start_matches('/'));
66            match tokio::fs::read(&file_path).await {
67                Ok(bytes) => {
68                    let mime = mime_from_path(&file_path);
69                    let mut headers = HeaderMap::new();
70                    headers.insert(header::CONTENT_TYPE, mime.parse().unwrap());
71                    headers.insert(header::CACHE_CONTROL, "no-cache".parse().unwrap());
72                    (StatusCode::OK, headers, bytes).into_response()
73                }
74                Err(_) => StatusCode::NOT_FOUND.into_response(),
75            }
76        }
77        None => StatusCode::NOT_FOUND.into_response(),
78    }
79}
80
81// ── Widget HTML ─────────────────────────────────────────
82
83/// Fetch widget HTML for a given widget name (used by resources/read interception).
84/// Asset URLs are made absolute so they resolve through the tunnel, not the sandbox origin.
85pub async fn fetch_widget_html(state: &ProxyState, widget_name: &str) -> Option<String> {
86    let html = match &state.widget_source {
87        Some(WidgetSource::Proxy(base_url)) => {
88            let url = format!(
89                "{}/src/{}/index.html",
90                base_url.trim_end_matches('/'),
91                widget_name
92            );
93            let resp = state
94                .upstream
95                .http_client
96                .get(&url)
97                .timeout(Duration::from_secs(10))
98                .send()
99                .await
100                .ok()?;
101            if !resp.status().is_success() {
102                return None;
103            }
104            resp.text().await.ok()?
105        }
106        Some(WidgetSource::Static(dir)) => {
107            let path = PathBuf::from(dir).join(format!("src/{widget_name}/index.html"));
108            tokio::fs::read_to_string(&path).await.ok()?
109        }
110        None => return None,
111    };
112
113    // Make absolute paths point to our tunnel, not the sandbox origin
114    let config = state.rewrite_config.read().await;
115    let proxy = config.proxy_url.trim_end_matches('/');
116    Some(rewrite_html_asset_urls(&html, proxy))
117}
118
119/// Serve raw widget HTML at `/widgets/{name}.html`.
120pub async fn serve_widget_html(state: &ProxyState, name: &str) -> Response {
121    let Some(html) = fetch_widget_html(state, name).await else {
122        return (StatusCode::NOT_FOUND, format!("Widget '{name}' not found")).into_response();
123    };
124
125    let mut headers = HeaderMap::new();
126    headers.insert(
127        header::CONTENT_TYPE,
128        "text/html; charset=utf-8".parse().unwrap(),
129    );
130    headers.insert(header::CACHE_CONTROL, "no-cache".parse().unwrap());
131    (StatusCode::OK, headers, html).into_response()
132}
133
134// ── Widget listing ──────────────────────────────────────
135
136/// JSON list of available widgets at `/widgets`.
137pub async fn list_widgets(state: &ProxyState) -> Response {
138    let names = discover_widget_names(state).await;
139    let body = serde_json::json!({
140        "widgets": names.iter().map(|n| {
141            serde_json::json!({
142                "name": n,
143                "url": format!("/widgets/{n}.html"),
144            })
145        }).collect::<Vec<_>>(),
146    });
147    let mut headers = HeaderMap::new();
148    headers.insert(header::CONTENT_TYPE, "application/json".parse().unwrap());
149    (
150        StatusCode::OK,
151        headers,
152        serde_json::to_string(&body).unwrap(),
153    )
154        .into_response()
155}
156
157// ── Widget discovery ────────────────────────────────────
158
159/// Discover available widget names from the widget source.
160pub async fn discover_widget_names(state: &ProxyState) -> Vec<String> {
161    match &state.widget_source {
162        Some(WidgetSource::Static(dir)) => {
163            let src_dir = PathBuf::from(dir).join("src");
164            let Ok(entries) = std::fs::read_dir(&src_dir) else {
165                return vec![];
166            };
167            let mut names: Vec<String> = entries
168                .filter_map(|e| e.ok())
169                .filter(|e| e.path().join("index.html").exists())
170                .filter_map(|e| e.file_name().into_string().ok())
171                .collect();
172            names.sort();
173            names
174        }
175        Some(WidgetSource::Proxy(base_url)) => {
176            let base = base_url.trim_end_matches('/');
177            let candidates = ["goal_detail", "question", "question_review", "vocab_review"];
178            let mut found = vec![];
179            for name in &candidates {
180                let url = format!("{base}/src/{name}/index.html");
181                if let Ok(resp) = state
182                    .upstream
183                    .http_client
184                    .head(&url)
185                    .timeout(Duration::from_secs(10))
186                    .send()
187                    .await
188                    && resp.status().is_success()
189                {
190                    found.push(name.to_string());
191                }
192            }
193            found
194        }
195        None => vec![],
196    }
197}
198
199// ── Helpers ─────────────────────────────────────────────
200
201fn mime_from_path(path: &Path) -> &'static str {
202    match path.extension().and_then(|e| e.to_str()) {
203        Some("js") => "application/javascript",
204        Some("css") => "text/css",
205        Some("html") => "text/html",
206        Some("svg") => "image/svg+xml",
207        Some("json") => "application/json",
208        Some("woff") => "font/woff",
209        Some("woff2") => "font/woff2",
210        Some("ttf") => "font/ttf",
211        Some("png") => "image/png",
212        Some("jpg" | "jpeg") => "image/jpeg",
213        _ => "application/octet-stream",
214    }
215}
216
217/// Rewrite asset URLs in HTML to point through the proxy.
218pub(crate) fn rewrite_html_asset_urls(html: &str, proxy_url: &str) -> String {
219    html.replace("\"/", &format!("\"{proxy_url}/"))
220        .replace("'/", &format!("'{proxy_url}/"))
221}
222
223#[cfg(test)]
224#[allow(non_snake_case)]
225mod tests {
226    use super::*;
227
228    // ── MIME type detection ──
229
230    #[test]
231    fn mime_from_path__js() {
232        assert_eq!(
233            mime_from_path(&PathBuf::from("app.js")),
234            "application/javascript"
235        );
236    }
237
238    #[test]
239    fn mime_from_path__css() {
240        assert_eq!(mime_from_path(&PathBuf::from("style.css")), "text/css");
241    }
242
243    #[test]
244    fn mime_from_path__html() {
245        assert_eq!(mime_from_path(&PathBuf::from("index.html")), "text/html");
246    }
247
248    #[test]
249    fn mime_from_path__svg() {
250        assert_eq!(mime_from_path(&PathBuf::from("icon.svg")), "image/svg+xml");
251    }
252
253    #[test]
254    fn mime_from_path__woff2() {
255        assert_eq!(mime_from_path(&PathBuf::from("font.woff2")), "font/woff2");
256    }
257
258    #[test]
259    fn mime_from_path__jpeg_variants() {
260        assert_eq!(mime_from_path(&PathBuf::from("photo.jpg")), "image/jpeg");
261        assert_eq!(mime_from_path(&PathBuf::from("photo.jpeg")), "image/jpeg");
262    }
263
264    #[test]
265    fn mime_from_path__unknown_extension() {
266        assert_eq!(
267            mime_from_path(&PathBuf::from("file.xyz")),
268            "application/octet-stream"
269        );
270    }
271
272    #[test]
273    fn mime_from_path__no_extension() {
274        assert_eq!(
275            mime_from_path(&PathBuf::from("Makefile")),
276            "application/octet-stream"
277        );
278    }
279
280    // ── HTML asset URL rewriting ──
281
282    #[test]
283    fn rewrite_html_asset_urls__double_quote_absolute() {
284        let html = r#"<script src="/assets/main.js"></script>"#;
285        let result = rewrite_html_asset_urls(html, "https://abc.tunnel.example.com");
286        assert_eq!(
287            result,
288            r#"<script src="https://abc.tunnel.example.com/assets/main.js"></script>"#
289        );
290    }
291
292    #[test]
293    fn rewrite_html_asset_urls__single_quote_absolute() {
294        let html = "<link href='/styles/app.css'>";
295        let result = rewrite_html_asset_urls(html, "https://abc.tunnel.example.com");
296        assert_eq!(
297            result,
298            "<link href='https://abc.tunnel.example.com/styles/app.css'>"
299        );
300    }
301
302    #[test]
303    fn rewrite_html_asset_urls__preserves_relative() {
304        let html = r#"<script src="./local.js"></script>"#;
305        let result = rewrite_html_asset_urls(html, "https://abc.tunnel.example.com");
306        assert_eq!(result, r#"<script src="./local.js"></script>"#);
307    }
308
309    #[test]
310    fn rewrite_html_asset_urls__preserves_external() {
311        let html = r#"<script src="https://cdn.example.com/lib.js"></script>"#;
312        let result = rewrite_html_asset_urls(html, "https://abc.tunnel.example.com");
313        assert_eq!(
314            result,
315            r#"<script src="https://cdn.example.com/lib.js"></script>"#
316        );
317    }
318
319    #[test]
320    fn rewrite_html_asset_urls__multiple_paths() {
321        let html = r#"<script src="/js/a.js"></script><link href="/css/b.css">"#;
322        let result = rewrite_html_asset_urls(html, "https://proxy.example.com");
323        assert!(result.contains("https://proxy.example.com/js/a.js"));
324        assert!(result.contains("https://proxy.example.com/css/b.css"));
325    }
326
327    #[test]
328    fn rewrite_html_asset_urls__strips_trailing_slash() {
329        let html = r#"<script src="/app.js"></script>"#;
330        let result = rewrite_html_asset_urls(html, "https://proxy.example.com");
331        assert!(result.contains("https://proxy.example.com/app.js"));
332        assert!(!result.contains("https://proxy.example.com//app.js"));
333    }
334}