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