Skip to main content

mcpr_core/proxy/
rewrite.rs

1//! # Response rewriting for widget CSP
2//!
3//! Mutates JSON-RPC response bodies so widgets see the right CSP regardless of
4//! which host renders them.
5//!
6//! Three things happen on every widget meta rewrite:
7//!
8//! 1. **Aggregate.** Upstream CSP domains are collected from both the OpenAI
9//!    (`openai/widgetCSP`) and spec (`ui.csp`) shapes, per directive.
10//! 2. **Merge.** [`super::csp::effective_domains`] applies the per-directive
11//!    mode, widget-scoped overrides, and proxy URL to produce one domain list
12//!    per directive.
13//! 3. **Emit both shapes.** The merge result is written to *both*
14//!    `openai/widgetCSP` (snake_case) and `ui.csp` (camelCase). ChatGPT reads
15//!    the former, Claude and VS Code read the latter; unknown keys are
16//!    ignored, so emitting both means the same declared config works on every
17//!    host.
18//!
19//! A deep scan walks the entire response afterwards and prepends the proxy URL
20//! to any CSP domain array it finds, catching servers that embed CSP in
21//! non-standard locations.
22//!
23//! Response body `text` and `blob` fields are never touched — widget HTML is
24//! served verbatim.
25
26use serde_json::Value;
27
28use super::csp::{CspConfig, Directive, effective_domains, is_public_proxy_origin};
29use crate::protocol as jsonrpc;
30
31/// Runtime configuration for response rewriting.
32#[derive(Clone)]
33pub struct RewriteConfig {
34    /// Proxy URL (scheme + host, no trailing slash) to insert into every CSP
35    /// array so widgets can reach the proxy.
36    pub proxy_url: String,
37    /// Bare proxy host (no scheme) used when rewriting `widgetDomain`.
38    pub proxy_domain: String,
39    /// Upstream MCP URL, used to recognise and strip upstream self-references
40    /// from the CSP arrays the server returns.
41    pub mcp_upstream: String,
42    /// Declarative CSP config — global policies plus widget-scoped overrides.
43    pub csp: CspConfig,
44}
45
46impl RewriteConfig {
47    /// Wrap this config in the lock-free `ArcSwap` shape expected by
48    /// [`crate::proxy::ProxyState::rewrite_config`]. Saves every caller
49    /// from importing `arc_swap` directly.
50    pub fn into_swap(self) -> std::sync::Arc<arc_swap::ArcSwap<RewriteConfig>> {
51        std::sync::Arc::new(arc_swap::ArcSwap::from_pointee(self))
52    }
53}
54
55/// Rewrite a JSON-RPC response in place for the given method.
56///
57/// Returns `true` iff the body was actually mutated. Callers use this to
58/// decide whether the response needs to be re-serialized — see
59/// `pipeline::handlers::buffered` for the buffered-path dispatcher.
60#[must_use]
61pub fn rewrite_response(method: &str, body: &mut Value, config: &RewriteConfig) -> bool {
62    let mut mutated = false;
63    match method {
64        jsonrpc::TOOLS_LIST => {
65            if let Some(tools) = body
66                .get_mut("result")
67                .and_then(|r| r.get_mut("tools"))
68                .and_then(|t| t.as_array_mut())
69            {
70                for tool in tools {
71                    if let Some(meta) = tool.get_mut("_meta") {
72                        rewrite_widget_meta(meta, None, config);
73                        mutated = true;
74                    }
75                }
76            }
77        }
78        jsonrpc::TOOLS_CALL => {
79            if let Some(meta) = body.get_mut("result").and_then(|r| r.get_mut("_meta")) {
80                rewrite_widget_meta(meta, None, config);
81                mutated = true;
82            }
83        }
84        jsonrpc::RESOURCES_LIST => {
85            if let Some(resources) = body
86                .get_mut("result")
87                .and_then(|r| r.get_mut("resources"))
88                .and_then(|r| r.as_array_mut())
89            {
90                for resource in resources {
91                    let uri = resource
92                        .get("uri")
93                        .and_then(|v| v.as_str())
94                        .map(String::from);
95                    // A URI alone is enough to treat this as a widget resource
96                    // per the MCP Apps spec, so synthesize `_meta` when the
97                    // upstream omits it — otherwise declared CSP silently
98                    // wouldn't apply to under-declaring servers.
99                    let has_existing_meta = resource.get("_meta").is_some();
100                    if (uri.is_some() || has_existing_meta)
101                        && let Some(meta) = ensure_meta(resource)
102                    {
103                        rewrite_widget_meta(meta, uri.as_deref(), config);
104                        mutated = true;
105                    }
106                }
107            }
108        }
109        jsonrpc::RESOURCES_TEMPLATES_LIST => {
110            if let Some(templates) = body
111                .get_mut("result")
112                .and_then(|r| r.get_mut("resourceTemplates"))
113                .and_then(|t| t.as_array_mut())
114            {
115                for template in templates {
116                    // Templates expose `uriTemplate`, not a concrete URI; treat
117                    // it as the match key so operators can glob on template IDs.
118                    let uri = template
119                        .get("uriTemplate")
120                        .and_then(|v| v.as_str())
121                        .map(String::from);
122                    let has_existing_meta = template.get("_meta").is_some();
123                    if (uri.is_some() || has_existing_meta)
124                        && let Some(meta) = ensure_meta(template)
125                    {
126                        rewrite_widget_meta(meta, uri.as_deref(), config);
127                        mutated = true;
128                    }
129                }
130            }
131        }
132        jsonrpc::RESOURCES_READ => {
133            if let Some(contents) = body
134                .get_mut("result")
135                .and_then(|r| r.get_mut("contents"))
136                .and_then(|c| c.as_array_mut())
137            {
138                for content in contents {
139                    let uri = content
140                        .get("uri")
141                        .and_then(|v| v.as_str())
142                        .map(String::from);
143                    let has_existing_meta = content.get("_meta").is_some();
144                    if (uri.is_some() || has_existing_meta)
145                        && let Some(meta) = ensure_meta(content)
146                    {
147                        rewrite_widget_meta(meta, uri.as_deref(), config);
148                        mutated = true;
149                    }
150                }
151            }
152        }
153        _ => {}
154    }
155
156    // Safety net: any CSP-shaped domain array elsewhere in the tree still gets
157    // the proxy URL. The merge rules above do not run here — this pass only
158    // guarantees the proxy URL is present.
159    mutated |= inject_proxy_into_all_csp(body, config);
160    mutated
161}
162
163/// Rewrite a widget metadata object.
164///
165/// Goal: an operator declares their CSP once in `mcpr.toml`, and the widget
166/// works on every host. Hosts disagree about where they read CSP from —
167/// ChatGPT reads `openai/widgetCSP` (snake_case), Claude and VS Code read
168/// `ui.csp` (camelCase). Instead of detecting the host, the rewrite does the
169/// same merge once and emits the result to *both* shapes. Extra keys are
170/// ignored by hosts that don't understand them.
171///
172/// Upstream domains that feed into the merge are aggregated from both shapes
173/// so a server declaring only one shape still informs the merge for the other.
174///
175/// `explicit_uri` is the resource URI the caller already resolved (for example,
176/// a `resources/read` caller knows the URI from the containing object). When
177/// missing, the URI is inferred from `meta.ui.resourceUri` or the legacy
178/// `openai/outputTemplate` field. The URI picks which `[[csp.widget]]`
179/// overrides apply.
180///
181/// The rewrite is skipped entirely for meta objects that show no sign of being
182/// a widget — a tool call result's `_meta`, for example — so non-widget metas
183/// are not polluted with CSP fields they don't need.
184fn rewrite_widget_meta(meta: &mut Value, explicit_uri: Option<&str>, config: &RewriteConfig) {
185    // Only overwrite when we actually have a public domain to write. If
186    // `proxy_domain` is empty (local-only dev with no `publicWidgetDomain`
187    // set), leave the upstream value alone rather than clobbering it with ""
188    // — and never synthesize the key when the upstream didn't declare it.
189    if !config.proxy_domain.is_empty() && meta.get("openai/widgetDomain").is_some() {
190        meta["openai/widgetDomain"] = Value::String(config.proxy_domain.clone());
191    }
192
193    if !is_widget_meta(meta, explicit_uri) {
194        // Inside rewrite_widget_meta: the caller's match arm has already
195        // flagged mutation, so the return value is uninteresting here.
196        let _ = inject_proxy_into_all_csp(meta, config);
197        return;
198    }
199
200    let inferred = explicit_uri
201        .map(String::from)
202        .or_else(|| extract_resource_uri(meta));
203    let uri = inferred.as_deref();
204    let upstream_host = strip_scheme(&config.mcp_upstream);
205
206    // Merge once per directive, using the union of upstream declarations from
207    // both shapes so a server that only declared `ui.csp` still informs the
208    // merge for the `openai/widgetCSP` output and vice versa.
209    let connect = merged_domains(meta, Directive::Connect, uri, &upstream_host, config);
210    let resource = merged_domains(meta, Directive::Resource, uri, &upstream_host, config);
211    let frame = merged_domains(meta, Directive::Frame, uri, &upstream_host, config);
212
213    write_openai_csp(meta, &connect, &resource, &frame);
214    write_spec_csp(meta, &connect, &resource, &frame);
215
216    // Same reasoning as above — caller already flags mutation.
217    let _ = inject_proxy_into_all_csp(meta, config);
218}
219
220/// Return `true` when the meta object belongs to a widget, either because it
221/// already holds widget-shaped fields or because the caller resolved an
222/// explicit resource URI for it.
223fn is_widget_meta(meta: &Value, explicit_uri: Option<&str>) -> bool {
224    if explicit_uri.is_some() {
225        return true;
226    }
227    meta.get("openai/widgetCSP").is_some()
228        || meta.get("openai/widgetDomain").is_some()
229        || meta.get("openai/outputTemplate").is_some()
230        || meta.pointer("/ui/csp").is_some()
231        || meta.pointer("/ui/resourceUri").is_some()
232        || meta.pointer("/ui/domain").is_some()
233}
234
235/// Extract a resource URI from a widget meta object. Prefers the spec field
236/// (`ui.resourceUri`) and falls back to the OpenAI legacy key.
237fn extract_resource_uri(meta: &Value) -> Option<String> {
238    if let Some(u) = meta.pointer("/ui/resourceUri").and_then(|v| v.as_str()) {
239        return Some(u.to_string());
240    }
241    meta.get("openai/outputTemplate")
242        .and_then(|v| v.as_str())
243        .map(String::from)
244}
245
246/// Compute the effective domain list for one directive, seeded from whatever
247/// the upstream server declared in either CSP shape.
248fn merged_domains(
249    meta: &Value,
250    directive: Directive,
251    resource_uri: Option<&str>,
252    upstream_host: &str,
253    config: &RewriteConfig,
254) -> Vec<String> {
255    let upstream = collect_upstream(meta, directive);
256    effective_domains(
257        &config.csp,
258        directive,
259        resource_uri,
260        &upstream,
261        upstream_host,
262        &config.proxy_url,
263    )
264}
265
266/// Gather every string domain the upstream declared for `directive`, looking
267/// at both `openai/widgetCSP` and `ui.csp`. Duplicates are removed in order.
268fn collect_upstream(meta: &Value, directive: Directive) -> Vec<String> {
269    let (openai_key, spec_key) = match directive {
270        Directive::Connect => ("connect_domains", "connectDomains"),
271        Directive::Resource => ("resource_domains", "resourceDomains"),
272        Directive::Frame => ("frame_domains", "frameDomains"),
273    };
274
275    let mut out: Vec<String> = Vec::new();
276    let mut append = |arr: &Vec<Value>| {
277        for v in arr {
278            if let Some(s) = v.as_str() {
279                let s = s.to_string();
280                if !out.contains(&s) {
281                    out.push(s);
282                }
283            }
284        }
285    };
286
287    if let Some(arr) = meta
288        .get("openai/widgetCSP")
289        .and_then(|c| c.get(openai_key))
290        .and_then(|v| v.as_array())
291    {
292        append(arr);
293    }
294    if let Some(arr) = meta
295        .pointer("/ui/csp")
296        .and_then(|c| c.get(spec_key))
297        .and_then(|v| v.as_array())
298    {
299        append(arr);
300    }
301    out
302}
303
304/// Write the OpenAI-shaped CSP block, creating the parent object when needed.
305fn write_openai_csp(meta: &mut Value, connect: &[String], resource: &[String], frame: &[String]) {
306    let Some(obj) = meta.as_object_mut() else {
307        return;
308    };
309    obj.insert(
310        "openai/widgetCSP".to_string(),
311        serde_json::json!({
312            "connect_domains": connect,
313            "resource_domains": resource,
314            "frame_domains": frame,
315        }),
316    );
317}
318
319/// Write the spec-shaped CSP block under `ui.csp`, creating `ui` when needed.
320fn write_spec_csp(meta: &mut Value, connect: &[String], resource: &[String], frame: &[String]) {
321    let Some(obj) = meta.as_object_mut() else {
322        return;
323    };
324    let ui = obj
325        .entry("ui".to_string())
326        .or_insert_with(|| Value::Object(serde_json::Map::new()));
327    if !ui.is_object() {
328        *ui = Value::Object(serde_json::Map::new());
329    }
330    let ui_obj = ui.as_object_mut().unwrap();
331    ui_obj.insert(
332        "csp".to_string(),
333        serde_json::json!({
334            "connectDomains": connect,
335            "resourceDomains": resource,
336            "frameDomains": frame,
337        }),
338    );
339}
340
341/// Recursively ensure the proxy URL is present in every CSP-shaped domain array
342/// that needs to reach the proxy (connect, resource). Frame arrays are skipped —
343/// see `csp::effective_domains`.
344///
345/// Does not apply the merge rules — that would require URI context the deep scan
346/// does not have. The only guarantee is "proxy URL is reachable from the widget."
347/// Walk `value` and insert `config.proxy_url` at the front of every CSP
348/// domain array that doesn't already contain it. Returns `true` iff any
349/// insertion happened — lets callers skip re-serialization on no-op walks.
350#[must_use]
351fn inject_proxy_into_all_csp(value: &mut Value, config: &RewriteConfig) -> bool {
352    // Skip entirely when there is no public origin worth injecting. A
353    // localhost proxy URL in a submitted widget's CSP is useless to the
354    // host (ChatGPT/Claude can't reach it) and just clutters the emitted
355    // config — better to leave the arrays alone.
356    if !is_public_proxy_origin(&config.proxy_url) {
357        return false;
358    }
359    let mut mutated = false;
360    match value {
361        Value::Object(map) => {
362            for key in [
363                "connect_domains",
364                "resource_domains",
365                "connectDomains",
366                "resourceDomains",
367            ] {
368                if let Some(arr) = map.get_mut(key).and_then(|v| v.as_array_mut()) {
369                    let has_proxy = arr.iter().any(|v| v.as_str() == Some(&config.proxy_url));
370                    if !has_proxy {
371                        arr.insert(0, Value::String(config.proxy_url.clone()));
372                        mutated = true;
373                    }
374                }
375            }
376            for (_, v) in map.iter_mut() {
377                mutated |= inject_proxy_into_all_csp(v, config);
378            }
379        }
380        Value::Array(arr) => {
381            for item in arr {
382                mutated |= inject_proxy_into_all_csp(item, config);
383            }
384        }
385        _ => {}
386    }
387    mutated
388}
389
390fn strip_scheme(url: &str) -> String {
391    url.trim_start_matches("https://")
392        .trim_start_matches("http://")
393        .split('/')
394        .next()
395        .unwrap_or("")
396        .to_string()
397}
398
399/// Get `container._meta`, inserting an empty object if absent. `None` only when
400/// the container isn't a JSON object — MCP resource/content entries always are,
401/// so callers can treat `None` as a malformed upstream and skip.
402fn ensure_meta(container: &mut Value) -> Option<&mut Value> {
403    let obj = container.as_object_mut()?;
404    Some(
405        obj.entry("_meta".to_string())
406            .or_insert_with(|| Value::Object(serde_json::Map::new())),
407    )
408}
409
410#[cfg(test)]
411#[allow(non_snake_case)]
412mod tests {
413    use super::*;
414    use crate::proxy::csp::{DirectivePolicy, Mode, WidgetScoped};
415    use serde_json::json;
416
417    // ── helpers ────────────────────────────────────────────────────────────
418
419    fn rewrite_config() -> RewriteConfig {
420        RewriteConfig {
421            proxy_url: "https://abc.tunnel.example.com".into(),
422            proxy_domain: "abc.tunnel.example.com".into(),
423            mcp_upstream: "http://localhost:9000".into(),
424            csp: CspConfig::default(),
425        }
426    }
427
428    fn as_strs(arr: &Value) -> Vec<&str> {
429        arr.as_array()
430            .unwrap()
431            .iter()
432            .map(|v| v.as_str().unwrap())
433            .collect()
434    }
435
436    // ── resources/read: HTML body is never touched ─────────────────────────
437
438    #[test]
439    fn rewrite_response__resources_read_preserves_html() {
440        let config = rewrite_config();
441        let mut body = json!({
442            "jsonrpc": "2.0", "id": 1,
443            "result": {
444                "contents": [{
445                    "uri": "ui://widget/question",
446                    "mimeType": "text/html",
447                    "text": "<html><script src=\"/assets/main.js\"></script></html>"
448                }]
449            }
450        });
451        let original = body["result"]["contents"][0]["text"]
452            .as_str()
453            .unwrap()
454            .to_string();
455
456        let _ = rewrite_response("resources/read", &mut body, &config);
457
458        assert_eq!(
459            body["result"]["contents"][0]["text"].as_str().unwrap(),
460            original
461        );
462    }
463
464    // ── resources/read: rewrites meta, not text ────────────────────────────
465
466    #[test]
467    fn rewrite_response__resources_read_rewrites_meta_not_text() {
468        let config = rewrite_config();
469        let mut body = json!({
470            "result": {
471                "contents": [{
472                    "uri": "ui://widget/question",
473                    "mimeType": "text/html",
474                    "text": "<html><body>Hello</body></html>",
475                    "_meta": {
476                        "openai/widgetDomain": "localhost:9000",
477                        "openai/widgetCSP": {
478                            "resource_domains": ["http://localhost:9000"],
479                            "connect_domains": ["http://localhost:9000"]
480                        }
481                    }
482                }]
483            }
484        });
485
486        let _ = rewrite_response("resources/read", &mut body, &config);
487
488        let content = &body["result"]["contents"][0];
489        assert_eq!(
490            content["text"].as_str().unwrap(),
491            "<html><body>Hello</body></html>"
492        );
493        assert_eq!(
494            content["_meta"]["openai/widgetDomain"].as_str().unwrap(),
495            "abc.tunnel.example.com"
496        );
497        let resources = as_strs(&content["_meta"]["openai/widgetCSP"]["resource_domains"]);
498        assert!(resources.contains(&"https://abc.tunnel.example.com"));
499        assert!(!resources.iter().any(|d| d.contains("localhost")));
500    }
501
502    // ── tools/list: per-tool meta rewrite ──────────────────────────────────
503
504    #[test]
505    fn rewrite_response__tools_list_rewrites_widget_domain() {
506        let config = rewrite_config();
507        let mut body = json!({
508            "result": {
509                "tools": [{
510                    "name": "create_question",
511                    "_meta": {
512                        "openai/widgetDomain": "old.domain.com",
513                        "openai/widgetCSP": {
514                            "resource_domains": ["http://localhost:4444"],
515                            "connect_domains": ["http://localhost:9000", "https://api.external.com"]
516                        }
517                    }
518                }]
519            }
520        });
521
522        let _ = rewrite_response("tools/list", &mut body, &config);
523
524        let meta = &body["result"]["tools"][0]["_meta"];
525        assert_eq!(
526            meta["openai/widgetDomain"].as_str().unwrap(),
527            "abc.tunnel.example.com"
528        );
529        let connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
530        assert!(connect.contains(&"https://abc.tunnel.example.com"));
531        assert!(connect.contains(&"https://api.external.com"));
532        assert!(!connect.iter().any(|d| d.contains("localhost")));
533    }
534
535    // ── tools/call: rewrites result.meta ───────────────────────────────────
536
537    #[test]
538    fn rewrite_response__tools_call_rewrites_meta() {
539        let config = rewrite_config();
540        let mut body = json!({
541            "result": {
542                "content": [{"type": "text", "text": "some result"}],
543                "_meta": {
544                    "openai/widgetDomain": "old.domain.com",
545                    "openai/widgetCSP": {
546                        "resource_domains": ["http://localhost:4444"]
547                    }
548                }
549            }
550        });
551
552        let _ = rewrite_response("tools/call", &mut body, &config);
553
554        assert_eq!(
555            body["result"]["_meta"]["openai/widgetDomain"]
556                .as_str()
557                .unwrap(),
558            "abc.tunnel.example.com"
559        );
560        assert_eq!(
561            body["result"]["content"][0]["text"].as_str().unwrap(),
562            "some result"
563        );
564    }
565
566    // ── resources/list ─────────────────────────────────────────────────────
567
568    #[test]
569    fn rewrite_response__resources_list_rewrites_meta() {
570        let config = rewrite_config();
571        let mut body = json!({
572            "result": {
573                "resources": [{
574                    "uri": "ui://widget/question",
575                    "name": "Question Widget",
576                    "_meta": {
577                        "openai/widgetDomain": "old.domain.com"
578                    }
579                }]
580            }
581        });
582
583        let _ = rewrite_response("resources/list", &mut body, &config);
584
585        assert_eq!(
586            body["result"]["resources"][0]["_meta"]["openai/widgetDomain"]
587                .as_str()
588                .unwrap(),
589            "abc.tunnel.example.com"
590        );
591    }
592
593    // ── resources/templates/list ───────────────────────────────────────────
594
595    #[test]
596    fn rewrite_response__resources_templates_list_rewrites_meta() {
597        let config = rewrite_config();
598        let mut body = json!({
599            "result": {
600                "resourceTemplates": [{
601                    "uriTemplate": "file:///{path}",
602                    "name": "File Access",
603                    "_meta": {
604                        "openai/widgetDomain": "old.domain.com",
605                        "openai/widgetCSP": {
606                            "resource_domains": ["http://localhost:4444"],
607                            "connect_domains": ["http://localhost:9000"]
608                        }
609                    }
610                }]
611            }
612        });
613
614        let _ = rewrite_response("resources/templates/list", &mut body, &config);
615
616        let meta = &body["result"]["resourceTemplates"][0]["_meta"];
617        assert_eq!(
618            meta["openai/widgetDomain"].as_str().unwrap(),
619            "abc.tunnel.example.com"
620        );
621        let resources = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
622        assert!(resources.contains(&"https://abc.tunnel.example.com"));
623        assert!(!resources.iter().any(|d| d.contains("localhost")));
624    }
625
626    // ── CSP merge: localhost stripping ─────────────────────────────────────
627
628    #[test]
629    fn rewrite_response__csp_strips_localhost() {
630        let config = rewrite_config();
631        let mut body = json!({
632            "result": {
633                "tools": [{
634                    "name": "test",
635                    "_meta": {
636                        "openai/widgetCSP": {
637                            "resource_domains": [
638                                "http://localhost:4444",
639                                "http://127.0.0.1:4444",
640                                "http://localhost:9000",
641                                "https://cdn.external.com"
642                            ]
643                        }
644                    }
645                }]
646            }
647        });
648
649        let _ = rewrite_response("tools/list", &mut body, &config);
650
651        let domains =
652            as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
653        assert_eq!(
654            domains,
655            vec!["https://abc.tunnel.example.com", "https://cdn.external.com"]
656        );
657    }
658
659    // ── CSP merge: declared global domains are appended ────────────────────
660
661    #[test]
662    fn rewrite_response__global_connect_domains_appended() {
663        let mut config = rewrite_config();
664        config.csp.connect_domains = DirectivePolicy {
665            domains: vec!["https://extra.example.com".into()],
666            mode: Mode::Extend,
667        };
668
669        let mut body = json!({
670            "result": {
671                "tools": [{
672                    "name": "test",
673                    "_meta": {
674                        "openai/widgetCSP": {
675                            "connect_domains": ["http://localhost:9000"]
676                        }
677                    }
678                }]
679            }
680        });
681
682        let _ = rewrite_response("tools/list", &mut body, &config);
683
684        let domains =
685            as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
686        assert!(domains.contains(&"https://extra.example.com"));
687        assert!(domains.contains(&"https://abc.tunnel.example.com"));
688    }
689
690    // ── CSP merge: no duplicate proxy entries ──────────────────────────────
691
692    #[test]
693    fn rewrite_response__csp_no_duplicate_proxy() {
694        let config = rewrite_config();
695        let mut body = json!({
696            "result": {
697                "tools": [{
698                    "name": "test",
699                    "_meta": {
700                        "openai/widgetCSP": {
701                            "resource_domains": ["https://abc.tunnel.example.com", "https://cdn.example.com"]
702                        }
703                    }
704                }]
705            }
706        });
707
708        let _ = rewrite_response("tools/list", &mut body, &config);
709
710        let domains =
711            as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
712        let count = domains
713            .iter()
714            .filter(|d| **d == "https://abc.tunnel.example.com")
715            .count();
716        assert_eq!(count, 1);
717    }
718
719    // ── Claude format parity ───────────────────────────────────────────────
720
721    #[test]
722    fn rewrite_response__claude_csp_format() {
723        let config = rewrite_config();
724        let mut body = json!({
725            "result": {
726                "tools": [{
727                    "name": "test",
728                    "_meta": {
729                        "ui": {
730                            "csp": {
731                                "connectDomains": ["http://localhost:9000"],
732                                "resourceDomains": ["http://localhost:4444"]
733                            }
734                        }
735                    }
736                }]
737            }
738        });
739
740        let _ = rewrite_response("tools/list", &mut body, &config);
741
742        let meta = &body["result"]["tools"][0]["_meta"]["ui"]["csp"];
743        let connect = as_strs(&meta["connectDomains"]);
744        let resource = as_strs(&meta["resourceDomains"]);
745        assert!(connect.contains(&"https://abc.tunnel.example.com"));
746        assert!(resource.contains(&"https://abc.tunnel.example.com"));
747        assert!(!connect.iter().any(|d| d.contains("localhost")));
748        assert!(!resource.iter().any(|d| d.contains("localhost")));
749    }
750
751    // ── Deep CSP injection fallback ────────────────────────────────────────
752
753    #[test]
754    fn rewrite_response__deep_csp_injection() {
755        let config = rewrite_config();
756        let mut body = json!({
757            "result": {
758                "content": [{
759                    "type": "text",
760                    "text": "result",
761                    "deeply": {
762                        "nested": {
763                            "connect_domains": ["https://only-external.com"]
764                        }
765                    }
766                }]
767            }
768        });
769
770        let _ = rewrite_response("tools/call", &mut body, &config);
771
772        let domains = as_strs(&body["result"]["content"][0]["deeply"]["nested"]["connect_domains"]);
773        assert!(domains.contains(&"https://abc.tunnel.example.com"));
774    }
775
776    #[test]
777    fn rewrite_response__deep_csp_injection_skips_frame_arrays() {
778        // Regression guard for the frame-domains fix: the safety-net deep scan
779        // must not prepend the proxy URL to frame arrays, or every mcpr-proxied
780        // widget would look like an iframe-embedder to ChatGPT and trigger
781        // extra security review.
782        let config = rewrite_config();
783        let mut body = json!({
784            "result": {
785                "content": [{
786                    "type": "text",
787                    "text": "result",
788                    "deeply": {
789                        "nested": {
790                            "frame_domains": ["https://embed.partner.com"],
791                            "frameDomains": ["https://embed.partner.com"]
792                        }
793                    }
794                }]
795            }
796        });
797
798        let _ = rewrite_response("tools/call", &mut body, &config);
799
800        let nested = &body["result"]["content"][0]["deeply"]["nested"];
801        let snake = as_strs(&nested["frame_domains"]);
802        let camel = as_strs(&nested["frameDomains"]);
803        assert_eq!(snake, vec!["https://embed.partner.com"]);
804        assert_eq!(camel, vec!["https://embed.partner.com"]);
805    }
806
807    // ── Unknown methods: only deep scan runs ───────────────────────────────
808
809    #[test]
810    fn rewrite_response__unknown_method_passthrough() {
811        let config = rewrite_config();
812        let mut body = json!({
813            "result": {
814                "data": "unchanged",
815                "_meta": { "openai/widgetDomain": "should-stay.com" }
816            }
817        });
818        let _ = rewrite_response("notifications/message", &mut body, &config);
819
820        assert_eq!(
821            body["result"]["_meta"]["openai/widgetDomain"]
822                .as_str()
823                .unwrap(),
824            "should-stay.com"
825        );
826        assert_eq!(body["result"]["data"].as_str().unwrap(), "unchanged");
827    }
828
829    // ── Global replace mode ───────────────────────────────────────────────
830
831    #[test]
832    fn rewrite_response__replace_mode_ignores_upstream() {
833        let mut config = rewrite_config();
834        config.csp.resource_domains = DirectivePolicy {
835            domains: vec!["https://allowed.example.com".into()],
836            mode: Mode::Replace,
837        };
838        config.csp.connect_domains = DirectivePolicy {
839            domains: vec!["https://allowed.example.com".into()],
840            mode: Mode::Replace,
841        };
842
843        let mut body = json!({
844            "result": {
845                "tools": [{
846                    "name": "test",
847                    "_meta": {
848                        "openai/widgetCSP": {
849                            "resource_domains": ["https://cdn.external.com", "https://api.external.com"],
850                            "connect_domains": ["https://api.external.com", "http://localhost:9000"]
851                        }
852                    }
853                }]
854            }
855        });
856
857        let _ = rewrite_response("tools/list", &mut body, &config);
858
859        let resources =
860            as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
861        assert_eq!(
862            resources,
863            vec![
864                "https://abc.tunnel.example.com",
865                "https://allowed.example.com"
866            ]
867        );
868        let connect =
869            as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
870        assert_eq!(
871            connect,
872            vec![
873                "https://abc.tunnel.example.com",
874                "https://allowed.example.com"
875            ]
876        );
877    }
878
879    // ── Widget-scoped overrides ────────────────────────────────────────────
880
881    #[test]
882    fn rewrite_response__widget_scope_matches_resource_uri() {
883        // A widget override should only apply when the resource URI matches.
884        let mut config = rewrite_config();
885        config.csp.widgets.push(WidgetScoped {
886            match_pattern: "ui://widget/payment*".into(),
887            connect_domains: vec!["https://api.stripe.com".into()],
888            connect_domains_mode: Mode::Extend,
889            ..Default::default()
890        });
891
892        let mut body = json!({
893            "result": {
894                "resources": [
895                    {
896                        "uri": "ui://widget/payment-form",
897                        "_meta": {
898                            "openai/widgetCSP": { "connect_domains": [] }
899                        }
900                    },
901                    {
902                        "uri": "ui://widget/search",
903                        "_meta": {
904                            "openai/widgetCSP": { "connect_domains": [] }
905                        }
906                    }
907                ]
908            }
909        });
910
911        let _ = rewrite_response("resources/list", &mut body, &config);
912
913        let payment_connect = as_strs(
914            &body["result"]["resources"][0]["_meta"]["openai/widgetCSP"]["connect_domains"],
915        );
916        assert!(payment_connect.contains(&"https://api.stripe.com"));
917
918        let search_connect = as_strs(
919            &body["result"]["resources"][1]["_meta"]["openai/widgetCSP"]["connect_domains"],
920        );
921        assert!(!search_connect.contains(&"https://api.stripe.com"));
922    }
923
924    #[test]
925    fn rewrite_response__widget_replace_mode_wipes_upstream() {
926        let mut config = rewrite_config();
927        config.csp.widgets.push(WidgetScoped {
928            match_pattern: "ui://widget/*".into(),
929            connect_domains: vec!["https://api.stripe.com".into()],
930            connect_domains_mode: Mode::Replace,
931            ..Default::default()
932        });
933
934        let mut body = json!({
935            "result": {
936                "contents": [{
937                    "uri": "ui://widget/payment",
938                    "_meta": {
939                        "openai/widgetCSP": {
940                            "connect_domains": [
941                                "https://api.external.com",
942                                "https://another.external.com"
943                            ]
944                        }
945                    }
946                }]
947            }
948        });
949
950        let _ = rewrite_response("resources/read", &mut body, &config);
951
952        let connect =
953            as_strs(&body["result"]["contents"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
954        assert_eq!(
955            connect,
956            vec!["https://abc.tunnel.example.com", "https://api.stripe.com"]
957        );
958    }
959
960    #[test]
961    fn rewrite_response__widget_uri_inferred_from_tool_meta() {
962        // tools/list responses do not carry a URI on the tool itself, but the
963        // widget resource URI lives in meta.ui.resourceUri; widget overrides
964        // should match against that.
965        let mut config = rewrite_config();
966        config.csp.widgets.push(WidgetScoped {
967            match_pattern: "ui://widget/payment*".into(),
968            connect_domains: vec!["https://api.stripe.com".into()],
969            connect_domains_mode: Mode::Extend,
970            ..Default::default()
971        });
972
973        let mut body = json!({
974            "result": {
975                "tools": [{
976                    "name": "take_payment",
977                    "_meta": {
978                        "ui": { "resourceUri": "ui://widget/payment-form" },
979                        "openai/widgetCSP": { "connect_domains": [] }
980                    }
981                }]
982            }
983        });
984
985        let _ = rewrite_response("tools/list", &mut body, &config);
986
987        let connect =
988            as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
989        assert!(connect.contains(&"https://api.stripe.com"));
990    }
991
992    // ── Both shapes emitted from one declaration ───────────────────────────
993
994    #[test]
995    fn rewrite_response__spec_only_upstream_also_emits_openai_shape() {
996        // Upstream only declared ui.csp (spec shape). The rewrite must also
997        // synthesize openai/widgetCSP so ChatGPT — which reads the legacy
998        // key — receives the same effective CSP.
999        let config = rewrite_config();
1000        let mut body = json!({
1001            "result": {
1002                "contents": [{
1003                    "uri": "ui://widget/search",
1004                    "mimeType": "text/html",
1005                    "_meta": {
1006                        "ui": {
1007                            "csp": {
1008                                "connectDomains": ["https://api.external.com"],
1009                                "resourceDomains": ["https://cdn.external.com"]
1010                            }
1011                        }
1012                    }
1013                }]
1014            }
1015        });
1016
1017        let _ = rewrite_response("resources/read", &mut body, &config);
1018
1019        let meta = &body["result"]["contents"][0]["_meta"];
1020        let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1021        let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1022        assert_eq!(oa_connect, spec_connect);
1023        assert!(oa_connect.contains(&"https://api.external.com"));
1024        assert!(oa_connect.contains(&"https://abc.tunnel.example.com"));
1025    }
1026
1027    #[test]
1028    fn rewrite_response__openai_only_upstream_also_emits_spec_shape() {
1029        // Reverse of the above: upstream only declared openai/widgetCSP and
1030        // Claude/VS Code clients must still see ui.csp.
1031        let config = rewrite_config();
1032        let mut body = json!({
1033            "result": {
1034                "contents": [{
1035                    "uri": "ui://widget/search",
1036                    "mimeType": "text/html",
1037                    "_meta": {
1038                        "openai/widgetCSP": {
1039                            "connect_domains": ["https://api.external.com"],
1040                            "resource_domains": ["https://cdn.external.com"]
1041                        }
1042                    }
1043                }]
1044            }
1045        });
1046
1047        let _ = rewrite_response("resources/read", &mut body, &config);
1048
1049        let meta = &body["result"]["contents"][0]["_meta"];
1050        let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1051        let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1052        assert_eq!(oa_connect, spec_connect);
1053        assert!(spec_connect.contains(&"https://api.external.com"));
1054        assert!(spec_connect.contains(&"https://abc.tunnel.example.com"));
1055    }
1056
1057    #[test]
1058    fn rewrite_response__declared_config_synthesizes_both_shapes_from_empty() {
1059        // The server declared neither CSP shape. The operator's mcpr.toml
1060        // declares connectDomains. Both shapes appear in the response with
1061        // the declared domain, keyed off the URI on the containing resource.
1062        let mut config = rewrite_config();
1063        config.csp.connect_domains = DirectivePolicy {
1064            domains: vec!["https://api.declared.com".into()],
1065            mode: Mode::Extend,
1066        };
1067
1068        let mut body = json!({
1069            "result": {
1070                "resources": [{
1071                    "uri": "ui://widget/search",
1072                    "_meta": {
1073                        "openai/widgetDomain": "old.domain.com"
1074                    }
1075                }]
1076            }
1077        });
1078
1079        let _ = rewrite_response("resources/list", &mut body, &config);
1080
1081        let meta = &body["result"]["resources"][0]["_meta"];
1082        let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1083        let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1084        assert_eq!(oa, spec);
1085        assert!(oa.contains(&"https://api.declared.com"));
1086        assert!(oa.contains(&"https://abc.tunnel.example.com"));
1087    }
1088
1089    #[test]
1090    fn rewrite_response__upstream_declarations_unioned_across_shapes() {
1091        // The server filled different domains into each shape. The merge must
1092        // see the union, not pick one.
1093        let config = rewrite_config();
1094        let mut body = json!({
1095            "result": {
1096                "contents": [{
1097                    "uri": "ui://widget/search",
1098                    "mimeType": "text/html",
1099                    "_meta": {
1100                        "openai/widgetCSP": {
1101                            "connect_domains": ["https://api.only-openai.com"]
1102                        },
1103                        "ui": {
1104                            "csp": {
1105                                "connectDomains": ["https://api.only-spec.com"]
1106                            }
1107                        }
1108                    }
1109                }]
1110            }
1111        });
1112
1113        let _ = rewrite_response("resources/read", &mut body, &config);
1114
1115        let meta = &body["result"]["contents"][0]["_meta"];
1116        let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1117        let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1118        assert_eq!(oa, spec);
1119        assert!(oa.contains(&"https://api.only-openai.com"));
1120        assert!(oa.contains(&"https://api.only-spec.com"));
1121    }
1122
1123    #[test]
1124    fn rewrite_response__non_widget_meta_is_not_polluted() {
1125        // A tool call result with plain meta (no widget indicators) must not
1126        // gain synthesized CSP fields — those only belong on widget metas.
1127        let config = rewrite_config();
1128        let mut body = json!({
1129            "result": {
1130                "content": [{"type": "text", "text": "plain result"}],
1131                "_meta": { "requestId": "abc-123" }
1132            }
1133        });
1134
1135        let _ = rewrite_response("tools/call", &mut body, &config);
1136
1137        let meta = &body["result"]["_meta"];
1138        assert!(meta.get("openai/widgetCSP").is_none());
1139        assert!(meta.get("ui").is_none());
1140        assert_eq!(meta["requestId"].as_str().unwrap(), "abc-123");
1141    }
1142
1143    #[test]
1144    fn rewrite_response__all_three_directives_synthesized() {
1145        // All three directives land in both shapes, not just connect.
1146        let mut config = rewrite_config();
1147        config.csp.connect_domains = DirectivePolicy {
1148            domains: vec!["https://api.example.com".into()],
1149            mode: Mode::Extend,
1150        };
1151        config.csp.resource_domains = DirectivePolicy {
1152            domains: vec!["https://cdn.example.com".into()],
1153            mode: Mode::Extend,
1154        };
1155
1156        let mut body = json!({
1157            "result": {
1158                "resources": [{
1159                    "uri": "ui://widget/search",
1160                    "_meta": { "openai/widgetDomain": "x" }
1161                }]
1162            }
1163        });
1164
1165        let _ = rewrite_response("resources/list", &mut body, &config);
1166
1167        let meta = &body["result"]["resources"][0]["_meta"];
1168        for shape in ["openai/widgetCSP"] {
1169            assert!(meta[shape]["connect_domains"].is_array());
1170            assert!(meta[shape]["resource_domains"].is_array());
1171            assert!(meta[shape]["frame_domains"].is_array());
1172        }
1173        assert!(meta["ui"]["csp"]["connectDomains"].is_array());
1174        assert!(meta["ui"]["csp"]["resourceDomains"].is_array());
1175        assert!(meta["ui"]["csp"]["frameDomains"].is_array());
1176    }
1177
1178    // ── Frame directive defaults strict ────────────────────────────────────
1179
1180    #[test]
1181    fn rewrite_response__frame_domains_default_replace_drops_upstream() {
1182        // Default config treats frameDomains as replace, so upstream values are
1183        // dropped. Unlike connect/resource, the proxy URL is NOT prepended to
1184        // frame — the widget doesn't iframe the proxy back into itself, and
1185        // prepending it flags the widget as an iframe-embedder to hosts.
1186        let config = rewrite_config();
1187        let mut body = json!({
1188            "result": {
1189                "tools": [{
1190                    "name": "test",
1191                    "_meta": {
1192                        "ui": {
1193                            "csp": {
1194                                "frameDomains": ["https://embed.external.com"]
1195                            }
1196                        }
1197                    }
1198                }]
1199            }
1200        });
1201
1202        let _ = rewrite_response("tools/list", &mut body, &config);
1203
1204        let frames = as_strs(&body["result"]["tools"][0]["_meta"]["ui"]["csp"]["frameDomains"]);
1205        assert!(
1206            frames.is_empty(),
1207            "frame_domains should be empty, got {frames:?}"
1208        );
1209    }
1210
1211    // ── End-to-end scenario ────────────────────────────────────────────────
1212
1213    #[test]
1214    fn rewrite_response__end_to_end_mcp_schema() {
1215        // Scenario exercising the full pipeline:
1216        // - A realistic multi-tool tools/list response from an upstream MCP
1217        //   server that declares widgets with mixed metadata shapes.
1218        // - A mcpr.toml with declared global CSP across all three directives
1219        //   and a widget-scoped override for the payment widget.
1220        //
1221        // After rewriting, every tool's meta should carry both CSP shapes
1222        // populated from the same merge, with upstream self-references
1223        // dropped, declared config applied, and the payment widget's Stripe
1224        // override only appearing on the payment tool.
1225        let mut config = rewrite_config();
1226        config.csp.connect_domains = DirectivePolicy {
1227            domains: vec!["https://api.myshop.com".into()],
1228            mode: Mode::Extend,
1229        };
1230        config.csp.resource_domains = DirectivePolicy {
1231            domains: vec!["https://cdn.myshop.com".into()],
1232            mode: Mode::Extend,
1233        };
1234        config.csp.widgets.push(WidgetScoped {
1235            match_pattern: "ui://widget/payment*".into(),
1236            connect_domains: vec!["https://api.stripe.com".into()],
1237            connect_domains_mode: Mode::Extend,
1238            resource_domains: vec!["https://js.stripe.com".into()],
1239            resource_domains_mode: Mode::Extend,
1240            ..Default::default()
1241        });
1242
1243        let mut body = json!({
1244            "jsonrpc": "2.0",
1245            "id": 42,
1246            "result": {
1247                "tools": [
1248                    {
1249                        "name": "search_products",
1250                        "description": "Search the product catalog",
1251                        "inputSchema": { "type": "object" },
1252                        "_meta": {
1253                            "openai/widgetDomain": "old.shop.com",
1254                            "openai/outputTemplate": "ui://widget/search",
1255                            "openai/widgetCSP": {
1256                                "connect_domains": ["http://localhost:9000"],
1257                                "resource_domains": ["http://localhost:4444"]
1258                            }
1259                        }
1260                    },
1261                    {
1262                        "name": "take_payment",
1263                        "description": "Charge a card",
1264                        "inputSchema": { "type": "object" },
1265                        "_meta": {
1266                            "ui": {
1267                                "resourceUri": "ui://widget/payment-form",
1268                                "csp": {
1269                                    "connectDomains": ["https://api.myshop.com"]
1270                                }
1271                            }
1272                        }
1273                    },
1274                    {
1275                        "name": "get_order_status",
1276                        "description": "Look up an order",
1277                        "inputSchema": { "type": "object" }
1278                    }
1279                ]
1280            }
1281        });
1282
1283        let _ = rewrite_response("tools/list", &mut body, &config);
1284
1285        let tools = body["result"]["tools"].as_array().unwrap();
1286
1287        // ── Tool 0: search — upstream declared only OpenAI shape ──────────
1288        let search_meta = &tools[0]["_meta"];
1289        assert_eq!(
1290            search_meta["openai/widgetDomain"].as_str().unwrap(),
1291            "abc.tunnel.example.com"
1292        );
1293        let search_oa_connect = as_strs(&search_meta["openai/widgetCSP"]["connect_domains"]);
1294        let search_spec_connect = as_strs(&search_meta["ui"]["csp"]["connectDomains"]);
1295        assert_eq!(search_oa_connect, search_spec_connect);
1296        // Proxy first, then declared global, upstream localhost dropped.
1297        assert_eq!(
1298            search_oa_connect,
1299            vec!["https://abc.tunnel.example.com", "https://api.myshop.com"]
1300        );
1301        // The payment widget override must NOT apply to the search tool.
1302        assert!(!search_oa_connect.contains(&"https://api.stripe.com"));
1303        // Frame directive defaults strict AND the proxy URL is not prepended
1304        // — so with no declared frame domains the array is fully empty.
1305        let search_oa_frame = as_strs(&search_meta["openai/widgetCSP"]["frame_domains"]);
1306        assert!(search_oa_frame.is_empty());
1307
1308        // ── Tool 1: payment — upstream declared only spec shape ──────────
1309        let payment_meta = &tools[1]["_meta"];
1310        let payment_oa_connect = as_strs(&payment_meta["openai/widgetCSP"]["connect_domains"]);
1311        let payment_spec_connect = as_strs(&payment_meta["ui"]["csp"]["connectDomains"]);
1312        assert_eq!(payment_oa_connect, payment_spec_connect);
1313        // Proxy + global + widget override (Stripe) all present, in that order.
1314        assert_eq!(
1315            payment_oa_connect,
1316            vec![
1317                "https://abc.tunnel.example.com",
1318                "https://api.myshop.com",
1319                "https://api.stripe.com",
1320            ]
1321        );
1322        let payment_oa_resource = as_strs(&payment_meta["openai/widgetCSP"]["resource_domains"]);
1323        assert_eq!(
1324            payment_oa_resource,
1325            vec![
1326                "https://abc.tunnel.example.com",
1327                "https://cdn.myshop.com",
1328                "https://js.stripe.com",
1329            ]
1330        );
1331
1332        // ── Tool 2: plain tool, no widget metadata ────────────────────────
1333        // Non-widget metas must not gain synthesized CSP fields.
1334        let plain = &tools[2];
1335        assert!(plain.get("_meta").is_none());
1336    }
1337
1338    // ── Real MCP wire shape: `_meta` key, not `meta` ──────────────────────
1339
1340    #[test]
1341    fn rewrite_response__tools_call_underscore_meta_is_rewritten() {
1342        // Regression: real MCP servers emit `_meta` (with underscore) per spec,
1343        // not `meta`. Earlier dispatch arms mistakenly read `meta`, silently
1344        // skipping every rewrite in production.
1345        let mut config = rewrite_config();
1346        config.csp.connect_domains = DirectivePolicy {
1347            domains: vec!["https://assets.usestudykit.com".into()],
1348            mode: Mode::Replace,
1349        };
1350        config.csp.resource_domains = DirectivePolicy {
1351            domains: vec!["https://assets.usestudykit.com".into()],
1352            mode: Mode::Replace,
1353        };
1354
1355        let mut body = json!({
1356            "result": {
1357                "_meta": {
1358                    "openai/outputTemplate": "ui://widget/vocab_review.html",
1359                    "openai/widgetDomain": "assets.usestudykit.com/src",
1360                    "openai/widgetCSP": {
1361                        "connect_domains": [
1362                            "http://localhost:9002",
1363                            "https://api.dictionaryapi.dev"
1364                        ],
1365                        "resource_domains": [
1366                            "http://localhost:9002",
1367                            "https://api.dictionaryapi.dev"
1368                        ]
1369                    },
1370                    "ui": {
1371                        "csp": {
1372                            "connectDomains": ["https://api.dictionaryapi.dev"],
1373                            "resourceDomains": ["https://api.dictionaryapi.dev"]
1374                        },
1375                        "resourceUri": "ui://widget/vocab_review.html"
1376                    }
1377                },
1378                "content": [{"type": "text", "text": "payload"}],
1379                "structuredContent": {"data": {"items": []}}
1380            }
1381        });
1382
1383        let _ = rewrite_response("tools/call", &mut body, &config);
1384
1385        let meta = &body["result"]["_meta"];
1386        assert_eq!(
1387            meta["openai/widgetDomain"].as_str().unwrap(),
1388            "abc.tunnel.example.com"
1389        );
1390        let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1391        let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1392        assert_eq!(oa_connect, spec_connect);
1393        assert_eq!(
1394            oa_connect,
1395            vec![
1396                "https://abc.tunnel.example.com",
1397                "https://assets.usestudykit.com"
1398            ]
1399        );
1400        let oa_resource = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
1401        assert_eq!(
1402            oa_resource,
1403            vec![
1404                "https://abc.tunnel.example.com",
1405                "https://assets.usestudykit.com"
1406            ]
1407        );
1408        assert_eq!(
1409            body["result"]["content"][0]["text"].as_str().unwrap(),
1410            "payload"
1411        );
1412    }
1413
1414    #[test]
1415    fn rewrite_response__resources_read_underscore_meta_is_rewritten() {
1416        let config = rewrite_config();
1417        let mut body = json!({
1418            "result": {
1419                "contents": [{
1420                    "uri": "ui://widget/question",
1421                    "mimeType": "text/html",
1422                    "text": "<html/>",
1423                    "_meta": {
1424                        "openai/widgetDomain": "old.domain.com"
1425                    }
1426                }]
1427            }
1428        });
1429
1430        let _ = rewrite_response("resources/read", &mut body, &config);
1431
1432        assert_eq!(
1433            body["result"]["contents"][0]["_meta"]["openai/widgetDomain"]
1434                .as_str()
1435                .unwrap(),
1436            "abc.tunnel.example.com"
1437        );
1438    }
1439
1440    #[test]
1441    fn rewrite_response__legacy_meta_key_is_ignored() {
1442        // Defensive: if an upstream sends the wrong key (`meta` without
1443        // underscore), we must not rewrite it — MCP spec uses `_meta` only.
1444        let config = rewrite_config();
1445        let mut body = json!({
1446            "result": {
1447                "_meta": {"openai/widgetDomain": "real.domain.com"},
1448                "meta":  {"openai/widgetDomain": "should-stay.com"}
1449            }
1450        });
1451
1452        let _ = rewrite_response("tools/call", &mut body, &config);
1453
1454        assert_eq!(
1455            body["result"]["_meta"]["openai/widgetDomain"]
1456                .as_str()
1457                .unwrap(),
1458            "abc.tunnel.example.com"
1459        );
1460        assert_eq!(
1461            body["result"]["meta"]["openai/widgetDomain"]
1462                .as_str()
1463                .unwrap(),
1464            "should-stay.com"
1465        );
1466    }
1467
1468    // ── _meta synthesis when upstream under-declares ──────────────────────
1469
1470    #[test]
1471    fn rewrite_response__resources_list_synthesizes_meta_when_upstream_omits() {
1472        // Widget resources with no `_meta` at all must still receive the
1473        // declared CSP — under-declaring servers are common in practice.
1474        let mut config = rewrite_config();
1475        config.csp.connect_domains = DirectivePolicy {
1476            domains: vec!["https://api.declared.com".into()],
1477            mode: Mode::Replace,
1478        };
1479        config.csp.resource_domains = DirectivePolicy {
1480            domains: vec!["https://cdn.declared.com".into()],
1481            mode: Mode::Extend,
1482        };
1483
1484        let mut body = json!({
1485            "result": {
1486                "resources": [{
1487                    "uri": "ui://widget/search",
1488                    "name": "Search Widget"
1489                }]
1490            }
1491        });
1492
1493        let mutated = rewrite_response("resources/list", &mut body, &config);
1494        assert!(mutated);
1495
1496        let meta = &body["result"]["resources"][0]["_meta"];
1497        let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1498        let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1499        assert_eq!(oa_connect, spec_connect);
1500        assert_eq!(
1501            oa_connect,
1502            vec!["https://abc.tunnel.example.com", "https://api.declared.com"]
1503        );
1504        let oa_resource = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
1505        assert!(oa_resource.contains(&"https://abc.tunnel.example.com"));
1506        assert!(oa_resource.contains(&"https://cdn.declared.com"));
1507    }
1508
1509    #[test]
1510    fn rewrite_response__resources_read_synthesizes_meta_when_upstream_omits() {
1511        let mut config = rewrite_config();
1512        config.csp.connect_domains = DirectivePolicy {
1513            domains: vec!["https://api.declared.com".into()],
1514            mode: Mode::Replace,
1515        };
1516
1517        let mut body = json!({
1518            "result": {
1519                "contents": [{
1520                    "uri": "ui://widget/question",
1521                    "mimeType": "text/html",
1522                    "text": "<html><body>Hello</body></html>"
1523                }]
1524            }
1525        });
1526
1527        let mutated = rewrite_response("resources/read", &mut body, &config);
1528        assert!(mutated);
1529
1530        assert_eq!(
1531            body["result"]["contents"][0]["text"].as_str().unwrap(),
1532            "<html><body>Hello</body></html>"
1533        );
1534        let meta = &body["result"]["contents"][0]["_meta"];
1535        let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1536        let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1537        assert_eq!(oa, spec);
1538        assert_eq!(
1539            oa,
1540            vec!["https://abc.tunnel.example.com", "https://api.declared.com"]
1541        );
1542    }
1543
1544    #[test]
1545    fn rewrite_response__resources_list_injects_into_empty_meta() {
1546        let mut config = rewrite_config();
1547        config.csp.connect_domains = DirectivePolicy {
1548            domains: vec!["https://api.declared.com".into()],
1549            mode: Mode::Extend,
1550        };
1551
1552        let mut body = json!({
1553            "result": {
1554                "resources": [{
1555                    "uri": "ui://widget/search",
1556                    "_meta": {}
1557                }]
1558            }
1559        });
1560
1561        let _ = rewrite_response("resources/list", &mut body, &config);
1562
1563        let meta = &body["result"]["resources"][0]["_meta"];
1564        let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1565        assert!(oa.contains(&"https://api.declared.com"));
1566        assert!(oa.contains(&"https://abc.tunnel.example.com"));
1567    }
1568
1569    #[test]
1570    fn rewrite_response__resources_templates_list_synthesizes_meta() {
1571        let mut config = rewrite_config();
1572        config.csp.resource_domains = DirectivePolicy {
1573            domains: vec!["https://cdn.declared.com".into()],
1574            mode: Mode::Extend,
1575        };
1576
1577        let mut body = json!({
1578            "result": {
1579                "resourceTemplates": [{
1580                    "uriTemplate": "ui://widget/{name}.html",
1581                    "name": "Widget Template"
1582                }]
1583            }
1584        });
1585
1586        let _ = rewrite_response("resources/templates/list", &mut body, &config);
1587
1588        let meta = &body["result"]["resourceTemplates"][0]["_meta"];
1589        let oa = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
1590        assert!(oa.contains(&"https://cdn.declared.com"));
1591        assert!(oa.contains(&"https://abc.tunnel.example.com"));
1592    }
1593
1594    #[test]
1595    fn rewrite_response__tools_call_no_meta_is_not_synthesized() {
1596        // Spec-aligned asymmetry: tools/call without widget indicators must
1597        // NOT grow synthesized CSP fields — CSP belongs on widget resources.
1598        let mut config = rewrite_config();
1599        config.csp.connect_domains = DirectivePolicy {
1600            domains: vec!["https://api.declared.com".into()],
1601            mode: Mode::Replace,
1602        };
1603
1604        let mut body = json!({
1605            "result": {
1606                "content": [{"type": "text", "text": "London 14C"}],
1607                "structuredContent": {"city": "London", "temp": 14}
1608            }
1609        });
1610
1611        let _ = rewrite_response("tools/call", &mut body, &config);
1612
1613        assert!(body["result"].get("_meta").is_none());
1614        assert_eq!(
1615            body["result"]["content"][0]["text"].as_str().unwrap(),
1616            "London 14C"
1617        );
1618    }
1619
1620    #[test]
1621    fn rewrite_response__resources_list_skips_when_no_uri_and_no_meta() {
1622        let config = rewrite_config();
1623        let mut body = json!({
1624            "result": {
1625                "resources": [{
1626                    "name": "malformed"
1627                }]
1628            }
1629        });
1630
1631        let _ = rewrite_response("resources/list", &mut body, &config);
1632
1633        assert!(body["result"]["resources"][0].get("_meta").is_none());
1634    }
1635
1636    // ── local-only mode: no public origin, don't pollute widget CSP ────────
1637
1638    fn local_only_config() -> RewriteConfig {
1639        // Mirrors what main.rs builds when tunnel is off and
1640        // `csp.publicWidgetDomain` is unset: proxy_url stays as the local
1641        // bind address for internal wiring, proxy_domain is empty to flag
1642        // "no public origin".
1643        RewriteConfig {
1644            proxy_url: "http://localhost:9002".into(),
1645            proxy_domain: String::new(),
1646            mcp_upstream: "http://localhost:9000".into(),
1647            csp: CspConfig::default(),
1648        }
1649    }
1650
1651    #[test]
1652    fn rewrite_response__local_only_leaves_widget_domain_untouched() {
1653        let config = local_only_config();
1654        let mut body = json!({
1655            "result": {
1656                "contents": [{
1657                    "uri": "ui://widget/card",
1658                    "_meta": {
1659                        "openai/widgetDomain": "dev.example.com"
1660                    }
1661                }]
1662            }
1663        });
1664
1665        let _ = rewrite_response("resources/read", &mut body, &config);
1666
1667        assert_eq!(
1668            body["result"]["contents"][0]["_meta"]["openai/widgetDomain"]
1669                .as_str()
1670                .unwrap(),
1671            "dev.example.com",
1672        );
1673    }
1674
1675    #[test]
1676    fn rewrite_response__local_only_skips_csp_injection() {
1677        let config = local_only_config();
1678        let mut body = json!({
1679            "result": {
1680                "contents": [{
1681                    "uri": "ui://widget/card",
1682                    "_meta": {
1683                        "openai/widgetCSP": {
1684                            "connect_domains": ["https://api.example.com"],
1685                            "resource_domains": ["https://cdn.example.com"],
1686                            "frame_domains": []
1687                        }
1688                    }
1689                }]
1690            }
1691        });
1692
1693        let _ = rewrite_response("resources/read", &mut body, &config);
1694
1695        let csp = &body["result"]["contents"][0]["_meta"]["openai/widgetCSP"];
1696        assert_eq!(
1697            as_strs(&csp["connect_domains"]),
1698            vec!["https://api.example.com"],
1699            "localhost proxy_url must not be injected",
1700        );
1701        assert_eq!(
1702            as_strs(&csp["resource_domains"]),
1703            vec!["https://cdn.example.com"],
1704        );
1705    }
1706
1707    #[test]
1708    fn rewrite_response__public_domain_is_injected() {
1709        // Sanity: with a public proxy origin, injection still happens.
1710        let config = rewrite_config();
1711        let mut body = json!({
1712            "result": {
1713                "contents": [{
1714                    "uri": "ui://widget/card",
1715                    "_meta": {
1716                        "openai/widgetCSP": {
1717                            "connect_domains": ["https://api.example.com"],
1718                            "resource_domains": [],
1719                            "frame_domains": []
1720                        }
1721                    }
1722                }]
1723            }
1724        });
1725
1726        let _ = rewrite_response("resources/read", &mut body, &config);
1727
1728        let connect =
1729            as_strs(&body["result"]["contents"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
1730        assert!(connect.contains(&"https://abc.tunnel.example.com"));
1731    }
1732}