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]
344fn inject_proxy_into_all_csp(value: &mut Value, config: &RewriteConfig) -> bool {
345 let mut mutated = false;
346 match value {
347 Value::Object(map) => {
348 for key in [
349 "connect_domains",
350 "resource_domains",
351 "frame_domains",
352 "connectDomains",
353 "resourceDomains",
354 "frameDomains",
355 ] {
356 if let Some(arr) = map.get_mut(key).and_then(|v| v.as_array_mut()) {
357 let has_proxy = arr.iter().any(|v| v.as_str() == Some(&config.proxy_url));
358 if !has_proxy {
359 arr.insert(0, Value::String(config.proxy_url.clone()));
360 mutated = true;
361 }
362 }
363 }
364 for (_, v) in map.iter_mut() {
365 mutated |= inject_proxy_into_all_csp(v, config);
366 }
367 }
368 Value::Array(arr) => {
369 for item in arr {
370 mutated |= inject_proxy_into_all_csp(item, config);
371 }
372 }
373 _ => {}
374 }
375 mutated
376}
377
378fn strip_scheme(url: &str) -> String {
379 url.trim_start_matches("https://")
380 .trim_start_matches("http://")
381 .split('/')
382 .next()
383 .unwrap_or("")
384 .to_string()
385}
386
387fn ensure_meta(container: &mut Value) -> Option<&mut Value> {
391 let obj = container.as_object_mut()?;
392 Some(
393 obj.entry("_meta".to_string())
394 .or_insert_with(|| Value::Object(serde_json::Map::new())),
395 )
396}
397
398#[cfg(test)]
399#[allow(non_snake_case)]
400mod tests {
401 use super::*;
402 use crate::proxy::csp::{DirectivePolicy, Mode, WidgetScoped};
403 use serde_json::json;
404
405 fn rewrite_config() -> RewriteConfig {
408 RewriteConfig {
409 proxy_url: "https://abc.tunnel.example.com".into(),
410 proxy_domain: "abc.tunnel.example.com".into(),
411 mcp_upstream: "http://localhost:9000".into(),
412 csp: CspConfig::default(),
413 }
414 }
415
416 fn as_strs(arr: &Value) -> Vec<&str> {
417 arr.as_array()
418 .unwrap()
419 .iter()
420 .map(|v| v.as_str().unwrap())
421 .collect()
422 }
423
424 #[test]
427 fn rewrite_response__resources_read_preserves_html() {
428 let config = rewrite_config();
429 let mut body = json!({
430 "jsonrpc": "2.0", "id": 1,
431 "result": {
432 "contents": [{
433 "uri": "ui://widget/question",
434 "mimeType": "text/html",
435 "text": "<html><script src=\"/assets/main.js\"></script></html>"
436 }]
437 }
438 });
439 let original = body["result"]["contents"][0]["text"]
440 .as_str()
441 .unwrap()
442 .to_string();
443
444 let _ = rewrite_response("resources/read", &mut body, &config);
445
446 assert_eq!(
447 body["result"]["contents"][0]["text"].as_str().unwrap(),
448 original
449 );
450 }
451
452 #[test]
455 fn rewrite_response__resources_read_rewrites_meta_not_text() {
456 let config = rewrite_config();
457 let mut body = json!({
458 "result": {
459 "contents": [{
460 "uri": "ui://widget/question",
461 "mimeType": "text/html",
462 "text": "<html><body>Hello</body></html>",
463 "_meta": {
464 "openai/widgetDomain": "localhost:9000",
465 "openai/widgetCSP": {
466 "resource_domains": ["http://localhost:9000"],
467 "connect_domains": ["http://localhost:9000"]
468 }
469 }
470 }]
471 }
472 });
473
474 let _ = rewrite_response("resources/read", &mut body, &config);
475
476 let content = &body["result"]["contents"][0];
477 assert_eq!(
478 content["text"].as_str().unwrap(),
479 "<html><body>Hello</body></html>"
480 );
481 assert_eq!(
482 content["_meta"]["openai/widgetDomain"].as_str().unwrap(),
483 "abc.tunnel.example.com"
484 );
485 let resources = as_strs(&content["_meta"]["openai/widgetCSP"]["resource_domains"]);
486 assert!(resources.contains(&"https://abc.tunnel.example.com"));
487 assert!(!resources.iter().any(|d| d.contains("localhost")));
488 }
489
490 #[test]
493 fn rewrite_response__tools_list_rewrites_widget_domain() {
494 let config = rewrite_config();
495 let mut body = json!({
496 "result": {
497 "tools": [{
498 "name": "create_question",
499 "_meta": {
500 "openai/widgetDomain": "old.domain.com",
501 "openai/widgetCSP": {
502 "resource_domains": ["http://localhost:4444"],
503 "connect_domains": ["http://localhost:9000", "https://api.external.com"]
504 }
505 }
506 }]
507 }
508 });
509
510 let _ = rewrite_response("tools/list", &mut body, &config);
511
512 let meta = &body["result"]["tools"][0]["_meta"];
513 assert_eq!(
514 meta["openai/widgetDomain"].as_str().unwrap(),
515 "abc.tunnel.example.com"
516 );
517 let connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
518 assert!(connect.contains(&"https://abc.tunnel.example.com"));
519 assert!(connect.contains(&"https://api.external.com"));
520 assert!(!connect.iter().any(|d| d.contains("localhost")));
521 }
522
523 #[test]
526 fn rewrite_response__tools_call_rewrites_meta() {
527 let config = rewrite_config();
528 let mut body = json!({
529 "result": {
530 "content": [{"type": "text", "text": "some result"}],
531 "_meta": {
532 "openai/widgetDomain": "old.domain.com",
533 "openai/widgetCSP": {
534 "resource_domains": ["http://localhost:4444"]
535 }
536 }
537 }
538 });
539
540 let _ = rewrite_response("tools/call", &mut body, &config);
541
542 assert_eq!(
543 body["result"]["_meta"]["openai/widgetDomain"]
544 .as_str()
545 .unwrap(),
546 "abc.tunnel.example.com"
547 );
548 assert_eq!(
549 body["result"]["content"][0]["text"].as_str().unwrap(),
550 "some result"
551 );
552 }
553
554 #[test]
557 fn rewrite_response__resources_list_rewrites_meta() {
558 let config = rewrite_config();
559 let mut body = json!({
560 "result": {
561 "resources": [{
562 "uri": "ui://widget/question",
563 "name": "Question Widget",
564 "_meta": {
565 "openai/widgetDomain": "old.domain.com"
566 }
567 }]
568 }
569 });
570
571 let _ = rewrite_response("resources/list", &mut body, &config);
572
573 assert_eq!(
574 body["result"]["resources"][0]["_meta"]["openai/widgetDomain"]
575 .as_str()
576 .unwrap(),
577 "abc.tunnel.example.com"
578 );
579 }
580
581 #[test]
584 fn rewrite_response__resources_templates_list_rewrites_meta() {
585 let config = rewrite_config();
586 let mut body = json!({
587 "result": {
588 "resourceTemplates": [{
589 "uriTemplate": "file:///{path}",
590 "name": "File Access",
591 "_meta": {
592 "openai/widgetDomain": "old.domain.com",
593 "openai/widgetCSP": {
594 "resource_domains": ["http://localhost:4444"],
595 "connect_domains": ["http://localhost:9000"]
596 }
597 }
598 }]
599 }
600 });
601
602 let _ = rewrite_response("resources/templates/list", &mut body, &config);
603
604 let meta = &body["result"]["resourceTemplates"][0]["_meta"];
605 assert_eq!(
606 meta["openai/widgetDomain"].as_str().unwrap(),
607 "abc.tunnel.example.com"
608 );
609 let resources = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
610 assert!(resources.contains(&"https://abc.tunnel.example.com"));
611 assert!(!resources.iter().any(|d| d.contains("localhost")));
612 }
613
614 #[test]
617 fn rewrite_response__csp_strips_localhost() {
618 let config = rewrite_config();
619 let mut body = json!({
620 "result": {
621 "tools": [{
622 "name": "test",
623 "_meta": {
624 "openai/widgetCSP": {
625 "resource_domains": [
626 "http://localhost:4444",
627 "http://127.0.0.1:4444",
628 "http://localhost:9000",
629 "https://cdn.external.com"
630 ]
631 }
632 }
633 }]
634 }
635 });
636
637 let _ = rewrite_response("tools/list", &mut body, &config);
638
639 let domains =
640 as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
641 assert_eq!(
642 domains,
643 vec!["https://abc.tunnel.example.com", "https://cdn.external.com"]
644 );
645 }
646
647 #[test]
650 fn rewrite_response__global_connect_domains_appended() {
651 let mut config = rewrite_config();
652 config.csp.connect_domains = DirectivePolicy {
653 domains: vec!["https://extra.example.com".into()],
654 mode: Mode::Extend,
655 };
656
657 let mut body = json!({
658 "result": {
659 "tools": [{
660 "name": "test",
661 "_meta": {
662 "openai/widgetCSP": {
663 "connect_domains": ["http://localhost:9000"]
664 }
665 }
666 }]
667 }
668 });
669
670 let _ = rewrite_response("tools/list", &mut body, &config);
671
672 let domains =
673 as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
674 assert!(domains.contains(&"https://extra.example.com"));
675 assert!(domains.contains(&"https://abc.tunnel.example.com"));
676 }
677
678 #[test]
681 fn rewrite_response__csp_no_duplicate_proxy() {
682 let config = rewrite_config();
683 let mut body = json!({
684 "result": {
685 "tools": [{
686 "name": "test",
687 "_meta": {
688 "openai/widgetCSP": {
689 "resource_domains": ["https://abc.tunnel.example.com", "https://cdn.example.com"]
690 }
691 }
692 }]
693 }
694 });
695
696 let _ = rewrite_response("tools/list", &mut body, &config);
697
698 let domains =
699 as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
700 let count = domains
701 .iter()
702 .filter(|d| **d == "https://abc.tunnel.example.com")
703 .count();
704 assert_eq!(count, 1);
705 }
706
707 #[test]
710 fn rewrite_response__claude_csp_format() {
711 let config = rewrite_config();
712 let mut body = json!({
713 "result": {
714 "tools": [{
715 "name": "test",
716 "_meta": {
717 "ui": {
718 "csp": {
719 "connectDomains": ["http://localhost:9000"],
720 "resourceDomains": ["http://localhost:4444"]
721 }
722 }
723 }
724 }]
725 }
726 });
727
728 let _ = rewrite_response("tools/list", &mut body, &config);
729
730 let meta = &body["result"]["tools"][0]["_meta"]["ui"]["csp"];
731 let connect = as_strs(&meta["connectDomains"]);
732 let resource = as_strs(&meta["resourceDomains"]);
733 assert!(connect.contains(&"https://abc.tunnel.example.com"));
734 assert!(resource.contains(&"https://abc.tunnel.example.com"));
735 assert!(!connect.iter().any(|d| d.contains("localhost")));
736 assert!(!resource.iter().any(|d| d.contains("localhost")));
737 }
738
739 #[test]
742 fn rewrite_response__deep_csp_injection() {
743 let config = rewrite_config();
744 let mut body = json!({
745 "result": {
746 "content": [{
747 "type": "text",
748 "text": "result",
749 "deeply": {
750 "nested": {
751 "connect_domains": ["https://only-external.com"]
752 }
753 }
754 }]
755 }
756 });
757
758 let _ = rewrite_response("tools/call", &mut body, &config);
759
760 let domains = as_strs(&body["result"]["content"][0]["deeply"]["nested"]["connect_domains"]);
761 assert!(domains.contains(&"https://abc.tunnel.example.com"));
762 }
763
764 #[test]
767 fn rewrite_response__unknown_method_passthrough() {
768 let config = rewrite_config();
769 let mut body = json!({
770 "result": {
771 "data": "unchanged",
772 "_meta": { "openai/widgetDomain": "should-stay.com" }
773 }
774 });
775 let _ = rewrite_response("notifications/message", &mut body, &config);
776
777 assert_eq!(
778 body["result"]["_meta"]["openai/widgetDomain"]
779 .as_str()
780 .unwrap(),
781 "should-stay.com"
782 );
783 assert_eq!(body["result"]["data"].as_str().unwrap(), "unchanged");
784 }
785
786 #[test]
789 fn rewrite_response__replace_mode_ignores_upstream() {
790 let mut config = rewrite_config();
791 config.csp.resource_domains = DirectivePolicy {
792 domains: vec!["https://allowed.example.com".into()],
793 mode: Mode::Replace,
794 };
795 config.csp.connect_domains = DirectivePolicy {
796 domains: vec!["https://allowed.example.com".into()],
797 mode: Mode::Replace,
798 };
799
800 let mut body = json!({
801 "result": {
802 "tools": [{
803 "name": "test",
804 "_meta": {
805 "openai/widgetCSP": {
806 "resource_domains": ["https://cdn.external.com", "https://api.external.com"],
807 "connect_domains": ["https://api.external.com", "http://localhost:9000"]
808 }
809 }
810 }]
811 }
812 });
813
814 let _ = rewrite_response("tools/list", &mut body, &config);
815
816 let resources =
817 as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
818 assert_eq!(
819 resources,
820 vec![
821 "https://abc.tunnel.example.com",
822 "https://allowed.example.com"
823 ]
824 );
825 let connect =
826 as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
827 assert_eq!(
828 connect,
829 vec![
830 "https://abc.tunnel.example.com",
831 "https://allowed.example.com"
832 ]
833 );
834 }
835
836 #[test]
839 fn rewrite_response__widget_scope_matches_resource_uri() {
840 let mut config = rewrite_config();
842 config.csp.widgets.push(WidgetScoped {
843 match_pattern: "ui://widget/payment*".into(),
844 connect_domains: vec!["https://api.stripe.com".into()],
845 connect_domains_mode: Mode::Extend,
846 ..Default::default()
847 });
848
849 let mut body = json!({
850 "result": {
851 "resources": [
852 {
853 "uri": "ui://widget/payment-form",
854 "_meta": {
855 "openai/widgetCSP": { "connect_domains": [] }
856 }
857 },
858 {
859 "uri": "ui://widget/search",
860 "_meta": {
861 "openai/widgetCSP": { "connect_domains": [] }
862 }
863 }
864 ]
865 }
866 });
867
868 let _ = rewrite_response("resources/list", &mut body, &config);
869
870 let payment_connect = as_strs(
871 &body["result"]["resources"][0]["_meta"]["openai/widgetCSP"]["connect_domains"],
872 );
873 assert!(payment_connect.contains(&"https://api.stripe.com"));
874
875 let search_connect = as_strs(
876 &body["result"]["resources"][1]["_meta"]["openai/widgetCSP"]["connect_domains"],
877 );
878 assert!(!search_connect.contains(&"https://api.stripe.com"));
879 }
880
881 #[test]
882 fn rewrite_response__widget_replace_mode_wipes_upstream() {
883 let mut config = rewrite_config();
884 config.csp.widgets.push(WidgetScoped {
885 match_pattern: "ui://widget/*".into(),
886 connect_domains: vec!["https://api.stripe.com".into()],
887 connect_domains_mode: Mode::Replace,
888 ..Default::default()
889 });
890
891 let mut body = json!({
892 "result": {
893 "contents": [{
894 "uri": "ui://widget/payment",
895 "_meta": {
896 "openai/widgetCSP": {
897 "connect_domains": [
898 "https://api.external.com",
899 "https://another.external.com"
900 ]
901 }
902 }
903 }]
904 }
905 });
906
907 let _ = rewrite_response("resources/read", &mut body, &config);
908
909 let connect =
910 as_strs(&body["result"]["contents"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
911 assert_eq!(
912 connect,
913 vec!["https://abc.tunnel.example.com", "https://api.stripe.com"]
914 );
915 }
916
917 #[test]
918 fn rewrite_response__widget_uri_inferred_from_tool_meta() {
919 let mut config = rewrite_config();
923 config.csp.widgets.push(WidgetScoped {
924 match_pattern: "ui://widget/payment*".into(),
925 connect_domains: vec!["https://api.stripe.com".into()],
926 connect_domains_mode: Mode::Extend,
927 ..Default::default()
928 });
929
930 let mut body = json!({
931 "result": {
932 "tools": [{
933 "name": "take_payment",
934 "_meta": {
935 "ui": { "resourceUri": "ui://widget/payment-form" },
936 "openai/widgetCSP": { "connect_domains": [] }
937 }
938 }]
939 }
940 });
941
942 let _ = rewrite_response("tools/list", &mut body, &config);
943
944 let connect =
945 as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
946 assert!(connect.contains(&"https://api.stripe.com"));
947 }
948
949 #[test]
952 fn rewrite_response__spec_only_upstream_also_emits_openai_shape() {
953 let config = rewrite_config();
957 let mut body = json!({
958 "result": {
959 "contents": [{
960 "uri": "ui://widget/search",
961 "mimeType": "text/html",
962 "_meta": {
963 "ui": {
964 "csp": {
965 "connectDomains": ["https://api.external.com"],
966 "resourceDomains": ["https://cdn.external.com"]
967 }
968 }
969 }
970 }]
971 }
972 });
973
974 let _ = rewrite_response("resources/read", &mut body, &config);
975
976 let meta = &body["result"]["contents"][0]["_meta"];
977 let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
978 let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
979 assert_eq!(oa_connect, spec_connect);
980 assert!(oa_connect.contains(&"https://api.external.com"));
981 assert!(oa_connect.contains(&"https://abc.tunnel.example.com"));
982 }
983
984 #[test]
985 fn rewrite_response__openai_only_upstream_also_emits_spec_shape() {
986 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 "openai/widgetCSP": {
996 "connect_domains": ["https://api.external.com"],
997 "resource_domains": ["https://cdn.external.com"]
998 }
999 }
1000 }]
1001 }
1002 });
1003
1004 let _ = rewrite_response("resources/read", &mut body, &config);
1005
1006 let meta = &body["result"]["contents"][0]["_meta"];
1007 let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1008 let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1009 assert_eq!(oa_connect, spec_connect);
1010 assert!(spec_connect.contains(&"https://api.external.com"));
1011 assert!(spec_connect.contains(&"https://abc.tunnel.example.com"));
1012 }
1013
1014 #[test]
1015 fn rewrite_response__declared_config_synthesizes_both_shapes_from_empty() {
1016 let mut config = rewrite_config();
1020 config.csp.connect_domains = DirectivePolicy {
1021 domains: vec!["https://api.declared.com".into()],
1022 mode: Mode::Extend,
1023 };
1024
1025 let mut body = json!({
1026 "result": {
1027 "resources": [{
1028 "uri": "ui://widget/search",
1029 "_meta": {
1030 "openai/widgetDomain": "old.domain.com"
1031 }
1032 }]
1033 }
1034 });
1035
1036 let _ = rewrite_response("resources/list", &mut body, &config);
1037
1038 let meta = &body["result"]["resources"][0]["_meta"];
1039 let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1040 let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1041 assert_eq!(oa, spec);
1042 assert!(oa.contains(&"https://api.declared.com"));
1043 assert!(oa.contains(&"https://abc.tunnel.example.com"));
1044 }
1045
1046 #[test]
1047 fn rewrite_response__upstream_declarations_unioned_across_shapes() {
1048 let config = rewrite_config();
1051 let mut body = json!({
1052 "result": {
1053 "contents": [{
1054 "uri": "ui://widget/search",
1055 "mimeType": "text/html",
1056 "_meta": {
1057 "openai/widgetCSP": {
1058 "connect_domains": ["https://api.only-openai.com"]
1059 },
1060 "ui": {
1061 "csp": {
1062 "connectDomains": ["https://api.only-spec.com"]
1063 }
1064 }
1065 }
1066 }]
1067 }
1068 });
1069
1070 let _ = rewrite_response("resources/read", &mut body, &config);
1071
1072 let meta = &body["result"]["contents"][0]["_meta"];
1073 let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1074 let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1075 assert_eq!(oa, spec);
1076 assert!(oa.contains(&"https://api.only-openai.com"));
1077 assert!(oa.contains(&"https://api.only-spec.com"));
1078 }
1079
1080 #[test]
1081 fn rewrite_response__non_widget_meta_is_not_polluted() {
1082 let config = rewrite_config();
1085 let mut body = json!({
1086 "result": {
1087 "content": [{"type": "text", "text": "plain result"}],
1088 "_meta": { "requestId": "abc-123" }
1089 }
1090 });
1091
1092 let _ = rewrite_response("tools/call", &mut body, &config);
1093
1094 let meta = &body["result"]["_meta"];
1095 assert!(meta.get("openai/widgetCSP").is_none());
1096 assert!(meta.get("ui").is_none());
1097 assert_eq!(meta["requestId"].as_str().unwrap(), "abc-123");
1098 }
1099
1100 #[test]
1101 fn rewrite_response__all_three_directives_synthesized() {
1102 let mut config = rewrite_config();
1104 config.csp.connect_domains = DirectivePolicy {
1105 domains: vec!["https://api.example.com".into()],
1106 mode: Mode::Extend,
1107 };
1108 config.csp.resource_domains = DirectivePolicy {
1109 domains: vec!["https://cdn.example.com".into()],
1110 mode: Mode::Extend,
1111 };
1112
1113 let mut body = json!({
1114 "result": {
1115 "resources": [{
1116 "uri": "ui://widget/search",
1117 "_meta": { "openai/widgetDomain": "x" }
1118 }]
1119 }
1120 });
1121
1122 let _ = rewrite_response("resources/list", &mut body, &config);
1123
1124 let meta = &body["result"]["resources"][0]["_meta"];
1125 for shape in ["openai/widgetCSP"] {
1126 assert!(meta[shape]["connect_domains"].is_array());
1127 assert!(meta[shape]["resource_domains"].is_array());
1128 assert!(meta[shape]["frame_domains"].is_array());
1129 }
1130 assert!(meta["ui"]["csp"]["connectDomains"].is_array());
1131 assert!(meta["ui"]["csp"]["resourceDomains"].is_array());
1132 assert!(meta["ui"]["csp"]["frameDomains"].is_array());
1133 }
1134
1135 #[test]
1138 fn rewrite_response__frame_domains_default_replace_drops_upstream() {
1139 let config = rewrite_config();
1142 let mut body = json!({
1143 "result": {
1144 "tools": [{
1145 "name": "test",
1146 "_meta": {
1147 "ui": {
1148 "csp": {
1149 "frameDomains": ["https://embed.external.com"]
1150 }
1151 }
1152 }
1153 }]
1154 }
1155 });
1156
1157 let _ = rewrite_response("tools/list", &mut body, &config);
1158
1159 let frames = as_strs(&body["result"]["tools"][0]["_meta"]["ui"]["csp"]["frameDomains"]);
1160 assert_eq!(frames, vec!["https://abc.tunnel.example.com"]);
1161 }
1162
1163 #[test]
1166 fn rewrite_response__end_to_end_mcp_schema() {
1167 let mut config = rewrite_config();
1178 config.csp.connect_domains = DirectivePolicy {
1179 domains: vec!["https://api.myshop.com".into()],
1180 mode: Mode::Extend,
1181 };
1182 config.csp.resource_domains = DirectivePolicy {
1183 domains: vec!["https://cdn.myshop.com".into()],
1184 mode: Mode::Extend,
1185 };
1186 config.csp.widgets.push(WidgetScoped {
1187 match_pattern: "ui://widget/payment*".into(),
1188 connect_domains: vec!["https://api.stripe.com".into()],
1189 connect_domains_mode: Mode::Extend,
1190 resource_domains: vec!["https://js.stripe.com".into()],
1191 resource_domains_mode: Mode::Extend,
1192 ..Default::default()
1193 });
1194
1195 let mut body = json!({
1196 "jsonrpc": "2.0",
1197 "id": 42,
1198 "result": {
1199 "tools": [
1200 {
1201 "name": "search_products",
1202 "description": "Search the product catalog",
1203 "inputSchema": { "type": "object" },
1204 "_meta": {
1205 "openai/widgetDomain": "old.shop.com",
1206 "openai/outputTemplate": "ui://widget/search",
1207 "openai/widgetCSP": {
1208 "connect_domains": ["http://localhost:9000"],
1209 "resource_domains": ["http://localhost:4444"]
1210 }
1211 }
1212 },
1213 {
1214 "name": "take_payment",
1215 "description": "Charge a card",
1216 "inputSchema": { "type": "object" },
1217 "_meta": {
1218 "ui": {
1219 "resourceUri": "ui://widget/payment-form",
1220 "csp": {
1221 "connectDomains": ["https://api.myshop.com"]
1222 }
1223 }
1224 }
1225 },
1226 {
1227 "name": "get_order_status",
1228 "description": "Look up an order",
1229 "inputSchema": { "type": "object" }
1230 }
1231 ]
1232 }
1233 });
1234
1235 let _ = rewrite_response("tools/list", &mut body, &config);
1236
1237 let tools = body["result"]["tools"].as_array().unwrap();
1238
1239 let search_meta = &tools[0]["_meta"];
1241 assert_eq!(
1242 search_meta["openai/widgetDomain"].as_str().unwrap(),
1243 "abc.tunnel.example.com"
1244 );
1245 let search_oa_connect = as_strs(&search_meta["openai/widgetCSP"]["connect_domains"]);
1246 let search_spec_connect = as_strs(&search_meta["ui"]["csp"]["connectDomains"]);
1247 assert_eq!(search_oa_connect, search_spec_connect);
1248 assert_eq!(
1250 search_oa_connect,
1251 vec!["https://abc.tunnel.example.com", "https://api.myshop.com"]
1252 );
1253 assert!(!search_oa_connect.contains(&"https://api.stripe.com"));
1255 let search_oa_frame = as_strs(&search_meta["openai/widgetCSP"]["frame_domains"]);
1257 assert_eq!(search_oa_frame, vec!["https://abc.tunnel.example.com"]);
1258
1259 let payment_meta = &tools[1]["_meta"];
1261 let payment_oa_connect = as_strs(&payment_meta["openai/widgetCSP"]["connect_domains"]);
1262 let payment_spec_connect = as_strs(&payment_meta["ui"]["csp"]["connectDomains"]);
1263 assert_eq!(payment_oa_connect, payment_spec_connect);
1264 assert_eq!(
1266 payment_oa_connect,
1267 vec![
1268 "https://abc.tunnel.example.com",
1269 "https://api.myshop.com",
1270 "https://api.stripe.com",
1271 ]
1272 );
1273 let payment_oa_resource = as_strs(&payment_meta["openai/widgetCSP"]["resource_domains"]);
1274 assert_eq!(
1275 payment_oa_resource,
1276 vec![
1277 "https://abc.tunnel.example.com",
1278 "https://cdn.myshop.com",
1279 "https://js.stripe.com",
1280 ]
1281 );
1282
1283 let plain = &tools[2];
1286 assert!(plain.get("_meta").is_none());
1287 }
1288
1289 #[test]
1292 fn rewrite_response__tools_call_underscore_meta_is_rewritten() {
1293 let mut config = rewrite_config();
1297 config.csp.connect_domains = DirectivePolicy {
1298 domains: vec!["https://assets.usestudykit.com".into()],
1299 mode: Mode::Replace,
1300 };
1301 config.csp.resource_domains = DirectivePolicy {
1302 domains: vec!["https://assets.usestudykit.com".into()],
1303 mode: Mode::Replace,
1304 };
1305
1306 let mut body = json!({
1307 "result": {
1308 "_meta": {
1309 "openai/outputTemplate": "ui://widget/vocab_review.html",
1310 "openai/widgetDomain": "assets.usestudykit.com/src",
1311 "openai/widgetCSP": {
1312 "connect_domains": [
1313 "http://localhost:9002",
1314 "https://api.dictionaryapi.dev"
1315 ],
1316 "resource_domains": [
1317 "http://localhost:9002",
1318 "https://api.dictionaryapi.dev"
1319 ]
1320 },
1321 "ui": {
1322 "csp": {
1323 "connectDomains": ["https://api.dictionaryapi.dev"],
1324 "resourceDomains": ["https://api.dictionaryapi.dev"]
1325 },
1326 "resourceUri": "ui://widget/vocab_review.html"
1327 }
1328 },
1329 "content": [{"type": "text", "text": "payload"}],
1330 "structuredContent": {"data": {"items": []}}
1331 }
1332 });
1333
1334 let _ = rewrite_response("tools/call", &mut body, &config);
1335
1336 let meta = &body["result"]["_meta"];
1337 assert_eq!(
1338 meta["openai/widgetDomain"].as_str().unwrap(),
1339 "abc.tunnel.example.com"
1340 );
1341 let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1342 let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1343 assert_eq!(oa_connect, spec_connect);
1344 assert_eq!(
1345 oa_connect,
1346 vec![
1347 "https://abc.tunnel.example.com",
1348 "https://assets.usestudykit.com"
1349 ]
1350 );
1351 let oa_resource = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
1352 assert_eq!(
1353 oa_resource,
1354 vec![
1355 "https://abc.tunnel.example.com",
1356 "https://assets.usestudykit.com"
1357 ]
1358 );
1359 assert_eq!(
1360 body["result"]["content"][0]["text"].as_str().unwrap(),
1361 "payload"
1362 );
1363 }
1364
1365 #[test]
1366 fn rewrite_response__resources_read_underscore_meta_is_rewritten() {
1367 let config = rewrite_config();
1368 let mut body = json!({
1369 "result": {
1370 "contents": [{
1371 "uri": "ui://widget/question",
1372 "mimeType": "text/html",
1373 "text": "<html/>",
1374 "_meta": {
1375 "openai/widgetDomain": "old.domain.com"
1376 }
1377 }]
1378 }
1379 });
1380
1381 let _ = rewrite_response("resources/read", &mut body, &config);
1382
1383 assert_eq!(
1384 body["result"]["contents"][0]["_meta"]["openai/widgetDomain"]
1385 .as_str()
1386 .unwrap(),
1387 "abc.tunnel.example.com"
1388 );
1389 }
1390
1391 #[test]
1392 fn rewrite_response__legacy_meta_key_is_ignored() {
1393 let config = rewrite_config();
1396 let mut body = json!({
1397 "result": {
1398 "_meta": {"openai/widgetDomain": "real.domain.com"},
1399 "meta": {"openai/widgetDomain": "should-stay.com"}
1400 }
1401 });
1402
1403 let _ = rewrite_response("tools/call", &mut body, &config);
1404
1405 assert_eq!(
1406 body["result"]["_meta"]["openai/widgetDomain"]
1407 .as_str()
1408 .unwrap(),
1409 "abc.tunnel.example.com"
1410 );
1411 assert_eq!(
1412 body["result"]["meta"]["openai/widgetDomain"]
1413 .as_str()
1414 .unwrap(),
1415 "should-stay.com"
1416 );
1417 }
1418
1419 #[test]
1422 fn rewrite_response__resources_list_synthesizes_meta_when_upstream_omits() {
1423 let mut config = rewrite_config();
1426 config.csp.connect_domains = DirectivePolicy {
1427 domains: vec!["https://api.declared.com".into()],
1428 mode: Mode::Replace,
1429 };
1430 config.csp.resource_domains = DirectivePolicy {
1431 domains: vec!["https://cdn.declared.com".into()],
1432 mode: Mode::Extend,
1433 };
1434
1435 let mut body = json!({
1436 "result": {
1437 "resources": [{
1438 "uri": "ui://widget/search",
1439 "name": "Search Widget"
1440 }]
1441 }
1442 });
1443
1444 let mutated = rewrite_response("resources/list", &mut body, &config);
1445 assert!(mutated);
1446
1447 let meta = &body["result"]["resources"][0]["_meta"];
1448 let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1449 let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1450 assert_eq!(oa_connect, spec_connect);
1451 assert_eq!(
1452 oa_connect,
1453 vec!["https://abc.tunnel.example.com", "https://api.declared.com"]
1454 );
1455 let oa_resource = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
1456 assert!(oa_resource.contains(&"https://abc.tunnel.example.com"));
1457 assert!(oa_resource.contains(&"https://cdn.declared.com"));
1458 }
1459
1460 #[test]
1461 fn rewrite_response__resources_read_synthesizes_meta_when_upstream_omits() {
1462 let mut config = rewrite_config();
1463 config.csp.connect_domains = DirectivePolicy {
1464 domains: vec!["https://api.declared.com".into()],
1465 mode: Mode::Replace,
1466 };
1467
1468 let mut body = json!({
1469 "result": {
1470 "contents": [{
1471 "uri": "ui://widget/question",
1472 "mimeType": "text/html",
1473 "text": "<html><body>Hello</body></html>"
1474 }]
1475 }
1476 });
1477
1478 let mutated = rewrite_response("resources/read", &mut body, &config);
1479 assert!(mutated);
1480
1481 assert_eq!(
1482 body["result"]["contents"][0]["text"].as_str().unwrap(),
1483 "<html><body>Hello</body></html>"
1484 );
1485 let meta = &body["result"]["contents"][0]["_meta"];
1486 let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1487 let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
1488 assert_eq!(oa, spec);
1489 assert_eq!(
1490 oa,
1491 vec!["https://abc.tunnel.example.com", "https://api.declared.com"]
1492 );
1493 }
1494
1495 #[test]
1496 fn rewrite_response__resources_list_injects_into_empty_meta() {
1497 let mut config = rewrite_config();
1498 config.csp.connect_domains = DirectivePolicy {
1499 domains: vec!["https://api.declared.com".into()],
1500 mode: Mode::Extend,
1501 };
1502
1503 let mut body = json!({
1504 "result": {
1505 "resources": [{
1506 "uri": "ui://widget/search",
1507 "_meta": {}
1508 }]
1509 }
1510 });
1511
1512 let _ = rewrite_response("resources/list", &mut body, &config);
1513
1514 let meta = &body["result"]["resources"][0]["_meta"];
1515 let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
1516 assert!(oa.contains(&"https://api.declared.com"));
1517 assert!(oa.contains(&"https://abc.tunnel.example.com"));
1518 }
1519
1520 #[test]
1521 fn rewrite_response__resources_templates_list_synthesizes_meta() {
1522 let mut config = rewrite_config();
1523 config.csp.resource_domains = DirectivePolicy {
1524 domains: vec!["https://cdn.declared.com".into()],
1525 mode: Mode::Extend,
1526 };
1527
1528 let mut body = json!({
1529 "result": {
1530 "resourceTemplates": [{
1531 "uriTemplate": "ui://widget/{name}.html",
1532 "name": "Widget Template"
1533 }]
1534 }
1535 });
1536
1537 let _ = rewrite_response("resources/templates/list", &mut body, &config);
1538
1539 let meta = &body["result"]["resourceTemplates"][0]["_meta"];
1540 let oa = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
1541 assert!(oa.contains(&"https://cdn.declared.com"));
1542 assert!(oa.contains(&"https://abc.tunnel.example.com"));
1543 }
1544
1545 #[test]
1546 fn rewrite_response__tools_call_no_meta_is_not_synthesized() {
1547 let mut config = rewrite_config();
1550 config.csp.connect_domains = DirectivePolicy {
1551 domains: vec!["https://api.declared.com".into()],
1552 mode: Mode::Replace,
1553 };
1554
1555 let mut body = json!({
1556 "result": {
1557 "content": [{"type": "text", "text": "London 14C"}],
1558 "structuredContent": {"city": "London", "temp": 14}
1559 }
1560 });
1561
1562 let _ = rewrite_response("tools/call", &mut body, &config);
1563
1564 assert!(body["result"].get("_meta").is_none());
1565 assert_eq!(
1566 body["result"]["content"][0]["text"].as_str().unwrap(),
1567 "London 14C"
1568 );
1569 }
1570
1571 #[test]
1572 fn rewrite_response__resources_list_skips_when_no_uri_and_no_meta() {
1573 let config = rewrite_config();
1574 let mut body = json!({
1575 "result": {
1576 "resources": [{
1577 "name": "malformed"
1578 }]
1579 }
1580 });
1581
1582 let _ = rewrite_response("resources/list", &mut body, &config);
1583
1584 assert!(body["result"]["resources"][0].get("_meta").is_none());
1585 }
1586}