1use serde_json::Value;
27
28use super::csp::{CspConfig, Directive, effective_domains};
29use crate::protocol as jsonrpc;
30
31#[derive(Clone)]
33pub struct RewriteConfig {
34 pub proxy_url: String,
37 pub proxy_domain: String,
39 pub mcp_upstream: String,
42 pub csp: CspConfig,
44}
45
46impl RewriteConfig {
47 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#[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 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 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 mutated |= inject_proxy_into_all_csp(body, config);
160 mutated
161}
162
163fn 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 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 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 let _ = inject_proxy_into_all_csp(meta, config);
214}
215
216fn 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
231fn 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
242fn 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
262fn 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
300fn 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
315fn 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#[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
388fn 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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[test]
871 fn rewrite_response__widget_scope_matches_resource_uri() {
872 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 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 #[test]
984 fn rewrite_response__spec_only_upstream_also_emits_openai_shape() {
985 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 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 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 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 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 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 #[test]
1170 fn rewrite_response__frame_domains_default_replace_drops_upstream() {
1171 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 #[test]
1203 fn rewrite_response__end_to_end_mcp_schema() {
1204 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 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 assert_eq!(
1287 search_oa_connect,
1288 vec!["https://abc.tunnel.example.com", "https://api.myshop.com"]
1289 );
1290 assert!(!search_oa_connect.contains(&"https://api.stripe.com"));
1292 let search_oa_frame = as_strs(&search_meta["openai/widgetCSP"]["frame_domains"]);
1295 assert!(search_oa_frame.is_empty());
1296
1297 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 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 let plain = &tools[2];
1324 assert!(plain.get("_meta").is_none());
1325 }
1326
1327 #[test]
1330 fn rewrite_response__tools_call_underscore_meta_is_rewritten() {
1331 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 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 #[test]
1460 fn rewrite_response__resources_list_synthesizes_meta_when_upstream_omits() {
1461 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 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}