mcpr_core/proxy/pipeline/middleware/
widget_overlay.rs

1//! `WidgetOverlayMiddleware` — substitute upstream-returned widget HTML with a local
2//! bundle for `ui://widget/*` resources. Runs only when:
3//!
4//! * the proxy has a widget source configured,
5//! * the request is a non-batch `resources/read`, and
6//! * the requested URI matches `ui://widget/<name>(.html)?`.
7//!
8//! The upstream response is still issued (handles auth, meta, CSP). This middleware
9//! only swaps the `text` field of each matching `contents[]` entry.
10
11use crate::protocol::McpMethod;
12use async_trait::async_trait;
13use serde_json::Value;
14
15use super::ResponseMiddleware;
16use crate::proxy::pipeline::context::{RequestContext, ResponseContext};
17use crate::proxy::proxy_state::ProxyState;
18use crate::proxy::widgets::fetch_widget_html;
19
20pub struct WidgetOverlayMiddleware;
21
22#[cfg(test)]
23#[allow(non_snake_case)]
24mod tests {
25    use std::sync::Arc;
26    use std::time::{Duration, Instant};
27
28    use axum::http::{HeaderMap, Method};
29    use serde_json::json;
30    use tokio::sync::RwLock;
31
32    use super::*;
33    use crate::protocol::schema_manager::{MemorySchemaStore, SchemaManager};
34    use crate::protocol::session::MemorySessionStore;
35    use crate::proxy::forwarding::UpstreamClient;
36    use crate::proxy::widgets::WidgetSource;
37    use crate::proxy::{CspConfig, RewriteConfig, new_shared_health};
38
39    fn proxy_with_widgets(dir: &std::path::Path) -> ProxyState {
40        ProxyState {
41            name: "t".into(),
42            mcp_upstream: "http://u".into(),
43            upstream: UpstreamClient {
44                http_client: reqwest::Client::builder().build().unwrap(),
45                semaphore: Arc::new(tokio::sync::Semaphore::new(1)),
46                request_timeout: Duration::from_secs(1),
47            },
48            max_request_body: 1024,
49            max_response_body: 1024,
50            rewrite_config: Arc::new(RwLock::new(RewriteConfig {
51                proxy_url: "http://p".into(),
52                proxy_domain: "p".into(),
53                mcp_upstream: "http://u".into(),
54                csp: CspConfig::default(),
55            })),
56            widget_source: Some(WidgetSource::Static(dir.to_string_lossy().to_string())),
57            sessions: MemorySessionStore::new(),
58            schema_manager: Arc::new(SchemaManager::new("t", MemorySchemaStore::new())),
59            health: new_shared_health(),
60            event_bus: crate::event::EventManager::new().start().bus,
61        }
62    }
63
64    fn batch_ctx_for_widget() -> RequestContext {
65        // Build a real batch ParsedBody so `is_batch` + `jsonrpc` stay
66        // consistent. The URI is what would otherwise match, proving that
67        // the batch flag alone gates the overlay.
68        let body = br#"[{"jsonrpc":"2.0","id":1,"method":"resources/read","params":{"uri":"ui://widget/question"}}]"#;
69        let parsed = crate::protocol::parse_body(body).unwrap();
70        assert!(parsed.is_batch);
71        RequestContext {
72            start: Instant::now(),
73            http_method: Method::POST,
74            path: "/mcp".into(),
75            request_size: body.len(),
76            wants_sse: false,
77            session_id: None,
78            jsonrpc: Some(parsed),
79            mcp_method: Some(McpMethod::ResourcesRead),
80            mcp_method_str: Some("resources/read".into()),
81            tool: None,
82            is_batch: true,
83            client_info_from_init: None,
84            client_name: None,
85            client_version: None,
86            tags: Vec::new(),
87        }
88    }
89
90    #[tokio::test]
91    async fn widget_overlay__batch_request_is_noop() {
92        let dir = tempfile::tempdir().unwrap();
93        let wdir = dir.path().join("src/question");
94        std::fs::create_dir_all(&wdir).unwrap();
95        std::fs::write(wdir.join("index.html"), "LOCAL").unwrap();
96
97        let state = proxy_with_widgets(dir.path());
98        let req = batch_ctx_for_widget();
99        let mut resp = ResponseContext::new(200, HeaderMap::new(), vec![], None);
100        resp.json = Some(json!({
101            "result": {"contents": [{"uri": "ui://widget/question", "text": "UPSTREAM"}]}
102        }));
103
104        WidgetOverlayMiddleware
105            .on_response(&state, &req, &mut resp)
106            .await;
107
108        let text = resp.json.as_ref().unwrap()["result"]["contents"][0]["text"]
109            .as_str()
110            .unwrap();
111        assert_eq!(text, "UPSTREAM", "batch requests must skip overlay");
112    }
113}
114
115#[async_trait]
116impl ResponseMiddleware for WidgetOverlayMiddleware {
117    async fn on_response(
118        &self,
119        state: &ProxyState,
120        req: &RequestContext,
121        resp: &mut ResponseContext,
122    ) {
123        if state.widget_source.is_none()
124            || req.is_batch
125            || req.mcp_method != Some(McpMethod::ResourcesRead)
126        {
127            return;
128        }
129        let Some(uri) = req
130            .jsonrpc
131            .as_ref()
132            .and_then(|p| p.first_params())
133            .and_then(|params| params.get("uri"))
134            .and_then(|u| u.as_str())
135            .and_then(|u| u.strip_prefix("ui://widget/"))
136            .map(|s| s.trim_end_matches(".html").to_string())
137        else {
138            return;
139        };
140
141        let Some(html) = fetch_widget_html(state, &uri).await else {
142            return;
143        };
144
145        let Some(json) = resp.json.as_mut() else {
146            return;
147        };
148        if let Some(contents) = json
149            .get_mut("result")
150            .and_then(|r| r.get_mut("contents"))
151            .and_then(|c| c.as_array_mut())
152        {
153            for content in contents.iter_mut() {
154                if content.get("text").is_some() {
155                    content["text"] = Value::String(html.clone());
156                }
157            }
158        }
159    }
160}
mcpr_core/proxy/pipeline/middleware/widget_overlay.rs

mcpr_core/proxy/pipeline/middleware/
widget_overlay.rs
