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