use serde_json::Value;
use super::csp::{CspConfig, Directive, effective_domains};
use crate::protocol as jsonrpc;
#[derive(Clone)]
pub struct RewriteConfig {
pub proxy_url: String,
pub proxy_domain: String,
pub mcp_upstream: String,
pub csp: CspConfig,
}
impl RewriteConfig {
pub fn into_swap(self) -> std::sync::Arc<arc_swap::ArcSwap<RewriteConfig>> {
std::sync::Arc::new(arc_swap::ArcSwap::from_pointee(self))
}
}
#[must_use]
pub fn rewrite_response(method: &str, body: &mut Value, config: &RewriteConfig) -> bool {
let mut mutated = false;
match method {
jsonrpc::TOOLS_LIST => {
if let Some(tools) = body
.get_mut("result")
.and_then(|r| r.get_mut("tools"))
.and_then(|t| t.as_array_mut())
{
for tool in tools {
if let Some(meta) = tool.get_mut("_meta") {
rewrite_widget_meta(meta, None, config);
mutated = true;
}
}
}
}
jsonrpc::TOOLS_CALL => {
if let Some(meta) = body.get_mut("result").and_then(|r| r.get_mut("_meta")) {
rewrite_widget_meta(meta, None, config);
mutated = true;
}
}
jsonrpc::RESOURCES_LIST => {
if let Some(resources) = body
.get_mut("result")
.and_then(|r| r.get_mut("resources"))
.and_then(|r| r.as_array_mut())
{
for resource in resources {
let uri = resource
.get("uri")
.and_then(|v| v.as_str())
.map(String::from);
let has_existing_meta = resource.get("_meta").is_some();
if (uri.is_some() || has_existing_meta)
&& let Some(meta) = ensure_meta(resource)
{
rewrite_widget_meta(meta, uri.as_deref(), config);
mutated = true;
}
}
}
}
jsonrpc::RESOURCES_TEMPLATES_LIST => {
if let Some(templates) = body
.get_mut("result")
.and_then(|r| r.get_mut("resourceTemplates"))
.and_then(|t| t.as_array_mut())
{
for template in templates {
let uri = template
.get("uriTemplate")
.and_then(|v| v.as_str())
.map(String::from);
let has_existing_meta = template.get("_meta").is_some();
if (uri.is_some() || has_existing_meta)
&& let Some(meta) = ensure_meta(template)
{
rewrite_widget_meta(meta, uri.as_deref(), config);
mutated = true;
}
}
}
}
jsonrpc::RESOURCES_READ => {
if let Some(contents) = body
.get_mut("result")
.and_then(|r| r.get_mut("contents"))
.and_then(|c| c.as_array_mut())
{
for content in contents {
let uri = content
.get("uri")
.and_then(|v| v.as_str())
.map(String::from);
let has_existing_meta = content.get("_meta").is_some();
if (uri.is_some() || has_existing_meta)
&& let Some(meta) = ensure_meta(content)
{
rewrite_widget_meta(meta, uri.as_deref(), config);
mutated = true;
}
}
}
}
_ => {}
}
mutated |= inject_proxy_into_all_csp(body, config);
mutated
}
fn rewrite_widget_meta(meta: &mut Value, explicit_uri: Option<&str>, config: &RewriteConfig) {
if meta.get("openai/widgetDomain").is_some() {
meta["openai/widgetDomain"] = Value::String(config.proxy_domain.clone());
}
if !is_widget_meta(meta, explicit_uri) {
let _ = inject_proxy_into_all_csp(meta, config);
return;
}
let inferred = explicit_uri
.map(String::from)
.or_else(|| extract_resource_uri(meta));
let uri = inferred.as_deref();
let upstream_host = strip_scheme(&config.mcp_upstream);
let connect = merged_domains(meta, Directive::Connect, uri, &upstream_host, config);
let resource = merged_domains(meta, Directive::Resource, uri, &upstream_host, config);
let frame = merged_domains(meta, Directive::Frame, uri, &upstream_host, config);
write_openai_csp(meta, &connect, &resource, &frame);
write_spec_csp(meta, &connect, &resource, &frame);
let _ = inject_proxy_into_all_csp(meta, config);
}
fn is_widget_meta(meta: &Value, explicit_uri: Option<&str>) -> bool {
if explicit_uri.is_some() {
return true;
}
meta.get("openai/widgetCSP").is_some()
|| meta.get("openai/widgetDomain").is_some()
|| meta.get("openai/outputTemplate").is_some()
|| meta.pointer("/ui/csp").is_some()
|| meta.pointer("/ui/resourceUri").is_some()
|| meta.pointer("/ui/domain").is_some()
}
fn extract_resource_uri(meta: &Value) -> Option<String> {
if let Some(u) = meta.pointer("/ui/resourceUri").and_then(|v| v.as_str()) {
return Some(u.to_string());
}
meta.get("openai/outputTemplate")
.and_then(|v| v.as_str())
.map(String::from)
}
fn merged_domains(
meta: &Value,
directive: Directive,
resource_uri: Option<&str>,
upstream_host: &str,
config: &RewriteConfig,
) -> Vec<String> {
let upstream = collect_upstream(meta, directive);
effective_domains(
&config.csp,
directive,
resource_uri,
&upstream,
upstream_host,
&config.proxy_url,
)
}
fn collect_upstream(meta: &Value, directive: Directive) -> Vec<String> {
let (openai_key, spec_key) = match directive {
Directive::Connect => ("connect_domains", "connectDomains"),
Directive::Resource => ("resource_domains", "resourceDomains"),
Directive::Frame => ("frame_domains", "frameDomains"),
};
let mut out: Vec<String> = Vec::new();
let mut append = |arr: &Vec<Value>| {
for v in arr {
if let Some(s) = v.as_str() {
let s = s.to_string();
if !out.contains(&s) {
out.push(s);
}
}
}
};
if let Some(arr) = meta
.get("openai/widgetCSP")
.and_then(|c| c.get(openai_key))
.and_then(|v| v.as_array())
{
append(arr);
}
if let Some(arr) = meta
.pointer("/ui/csp")
.and_then(|c| c.get(spec_key))
.and_then(|v| v.as_array())
{
append(arr);
}
out
}
fn write_openai_csp(meta: &mut Value, connect: &[String], resource: &[String], frame: &[String]) {
let Some(obj) = meta.as_object_mut() else {
return;
};
obj.insert(
"openai/widgetCSP".to_string(),
serde_json::json!({
"connect_domains": connect,
"resource_domains": resource,
"frame_domains": frame,
}),
);
}
fn write_spec_csp(meta: &mut Value, connect: &[String], resource: &[String], frame: &[String]) {
let Some(obj) = meta.as_object_mut() else {
return;
};
let ui = obj
.entry("ui".to_string())
.or_insert_with(|| Value::Object(serde_json::Map::new()));
if !ui.is_object() {
*ui = Value::Object(serde_json::Map::new());
}
let ui_obj = ui.as_object_mut().unwrap();
ui_obj.insert(
"csp".to_string(),
serde_json::json!({
"connectDomains": connect,
"resourceDomains": resource,
"frameDomains": frame,
}),
);
}
#[must_use]
fn inject_proxy_into_all_csp(value: &mut Value, config: &RewriteConfig) -> bool {
let mut mutated = false;
match value {
Value::Object(map) => {
for key in [
"connect_domains",
"resource_domains",
"frame_domains",
"connectDomains",
"resourceDomains",
"frameDomains",
] {
if let Some(arr) = map.get_mut(key).and_then(|v| v.as_array_mut()) {
let has_proxy = arr.iter().any(|v| v.as_str() == Some(&config.proxy_url));
if !has_proxy {
arr.insert(0, Value::String(config.proxy_url.clone()));
mutated = true;
}
}
}
for (_, v) in map.iter_mut() {
mutated |= inject_proxy_into_all_csp(v, config);
}
}
Value::Array(arr) => {
for item in arr {
mutated |= inject_proxy_into_all_csp(item, config);
}
}
_ => {}
}
mutated
}
fn strip_scheme(url: &str) -> String {
url.trim_start_matches("https://")
.trim_start_matches("http://")
.split('/')
.next()
.unwrap_or("")
.to_string()
}
fn ensure_meta(container: &mut Value) -> Option<&mut Value> {
let obj = container.as_object_mut()?;
Some(
obj.entry("_meta".to_string())
.or_insert_with(|| Value::Object(serde_json::Map::new())),
)
}
#[cfg(test)]
#[allow(non_snake_case)]
mod tests {
use super::*;
use crate::proxy::csp::{DirectivePolicy, Mode, WidgetScoped};
use serde_json::json;
fn rewrite_config() -> RewriteConfig {
RewriteConfig {
proxy_url: "https://abc.tunnel.example.com".into(),
proxy_domain: "abc.tunnel.example.com".into(),
mcp_upstream: "http://localhost:9000".into(),
csp: CspConfig::default(),
}
}
fn as_strs(arr: &Value) -> Vec<&str> {
arr.as_array()
.unwrap()
.iter()
.map(|v| v.as_str().unwrap())
.collect()
}
#[test]
fn rewrite_response__resources_read_preserves_html() {
let config = rewrite_config();
let mut body = json!({
"jsonrpc": "2.0", "id": 1,
"result": {
"contents": [{
"uri": "ui://widget/question",
"mimeType": "text/html",
"text": "<html><script src=\"/assets/main.js\"></script></html>"
}]
}
});
let original = body["result"]["contents"][0]["text"]
.as_str()
.unwrap()
.to_string();
let _ = rewrite_response("resources/read", &mut body, &config);
assert_eq!(
body["result"]["contents"][0]["text"].as_str().unwrap(),
original
);
}
#[test]
fn rewrite_response__resources_read_rewrites_meta_not_text() {
let config = rewrite_config();
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/question",
"mimeType": "text/html",
"text": "<html><body>Hello</body></html>",
"_meta": {
"openai/widgetDomain": "localhost:9000",
"openai/widgetCSP": {
"resource_domains": ["http://localhost:9000"],
"connect_domains": ["http://localhost:9000"]
}
}
}]
}
});
let _ = rewrite_response("resources/read", &mut body, &config);
let content = &body["result"]["contents"][0];
assert_eq!(
content["text"].as_str().unwrap(),
"<html><body>Hello</body></html>"
);
assert_eq!(
content["_meta"]["openai/widgetDomain"].as_str().unwrap(),
"abc.tunnel.example.com"
);
let resources = as_strs(&content["_meta"]["openai/widgetCSP"]["resource_domains"]);
assert!(resources.contains(&"https://abc.tunnel.example.com"));
assert!(!resources.iter().any(|d| d.contains("localhost")));
}
#[test]
fn rewrite_response__tools_list_rewrites_widget_domain() {
let config = rewrite_config();
let mut body = json!({
"result": {
"tools": [{
"name": "create_question",
"_meta": {
"openai/widgetDomain": "old.domain.com",
"openai/widgetCSP": {
"resource_domains": ["http://localhost:4444"],
"connect_domains": ["http://localhost:9000", "https://api.external.com"]
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let meta = &body["result"]["tools"][0]["_meta"];
assert_eq!(
meta["openai/widgetDomain"].as_str().unwrap(),
"abc.tunnel.example.com"
);
let connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
assert!(connect.contains(&"https://abc.tunnel.example.com"));
assert!(connect.contains(&"https://api.external.com"));
assert!(!connect.iter().any(|d| d.contains("localhost")));
}
#[test]
fn rewrite_response__tools_call_rewrites_meta() {
let config = rewrite_config();
let mut body = json!({
"result": {
"content": [{"type": "text", "text": "some result"}],
"_meta": {
"openai/widgetDomain": "old.domain.com",
"openai/widgetCSP": {
"resource_domains": ["http://localhost:4444"]
}
}
}
});
let _ = rewrite_response("tools/call", &mut body, &config);
assert_eq!(
body["result"]["_meta"]["openai/widgetDomain"]
.as_str()
.unwrap(),
"abc.tunnel.example.com"
);
assert_eq!(
body["result"]["content"][0]["text"].as_str().unwrap(),
"some result"
);
}
#[test]
fn rewrite_response__resources_list_rewrites_meta() {
let config = rewrite_config();
let mut body = json!({
"result": {
"resources": [{
"uri": "ui://widget/question",
"name": "Question Widget",
"_meta": {
"openai/widgetDomain": "old.domain.com"
}
}]
}
});
let _ = rewrite_response("resources/list", &mut body, &config);
assert_eq!(
body["result"]["resources"][0]["_meta"]["openai/widgetDomain"]
.as_str()
.unwrap(),
"abc.tunnel.example.com"
);
}
#[test]
fn rewrite_response__resources_templates_list_rewrites_meta() {
let config = rewrite_config();
let mut body = json!({
"result": {
"resourceTemplates": [{
"uriTemplate": "file:///{path}",
"name": "File Access",
"_meta": {
"openai/widgetDomain": "old.domain.com",
"openai/widgetCSP": {
"resource_domains": ["http://localhost:4444"],
"connect_domains": ["http://localhost:9000"]
}
}
}]
}
});
let _ = rewrite_response("resources/templates/list", &mut body, &config);
let meta = &body["result"]["resourceTemplates"][0]["_meta"];
assert_eq!(
meta["openai/widgetDomain"].as_str().unwrap(),
"abc.tunnel.example.com"
);
let resources = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
assert!(resources.contains(&"https://abc.tunnel.example.com"));
assert!(!resources.iter().any(|d| d.contains("localhost")));
}
#[test]
fn rewrite_response__csp_strips_localhost() {
let config = rewrite_config();
let mut body = json!({
"result": {
"tools": [{
"name": "test",
"_meta": {
"openai/widgetCSP": {
"resource_domains": [
"http://localhost:4444",
"http://127.0.0.1:4444",
"http://localhost:9000",
"https://cdn.external.com"
]
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let domains =
as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
assert_eq!(
domains,
vec!["https://abc.tunnel.example.com", "https://cdn.external.com"]
);
}
#[test]
fn rewrite_response__global_connect_domains_appended() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://extra.example.com".into()],
mode: Mode::Extend,
};
let mut body = json!({
"result": {
"tools": [{
"name": "test",
"_meta": {
"openai/widgetCSP": {
"connect_domains": ["http://localhost:9000"]
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let domains =
as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
assert!(domains.contains(&"https://extra.example.com"));
assert!(domains.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__csp_no_duplicate_proxy() {
let config = rewrite_config();
let mut body = json!({
"result": {
"tools": [{
"name": "test",
"_meta": {
"openai/widgetCSP": {
"resource_domains": ["https://abc.tunnel.example.com", "https://cdn.example.com"]
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let domains =
as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
let count = domains
.iter()
.filter(|d| **d == "https://abc.tunnel.example.com")
.count();
assert_eq!(count, 1);
}
#[test]
fn rewrite_response__claude_csp_format() {
let config = rewrite_config();
let mut body = json!({
"result": {
"tools": [{
"name": "test",
"_meta": {
"ui": {
"csp": {
"connectDomains": ["http://localhost:9000"],
"resourceDomains": ["http://localhost:4444"]
}
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let meta = &body["result"]["tools"][0]["_meta"]["ui"]["csp"];
let connect = as_strs(&meta["connectDomains"]);
let resource = as_strs(&meta["resourceDomains"]);
assert!(connect.contains(&"https://abc.tunnel.example.com"));
assert!(resource.contains(&"https://abc.tunnel.example.com"));
assert!(!connect.iter().any(|d| d.contains("localhost")));
assert!(!resource.iter().any(|d| d.contains("localhost")));
}
#[test]
fn rewrite_response__deep_csp_injection() {
let config = rewrite_config();
let mut body = json!({
"result": {
"content": [{
"type": "text",
"text": "result",
"deeply": {
"nested": {
"connect_domains": ["https://only-external.com"]
}
}
}]
}
});
let _ = rewrite_response("tools/call", &mut body, &config);
let domains = as_strs(&body["result"]["content"][0]["deeply"]["nested"]["connect_domains"]);
assert!(domains.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__unknown_method_passthrough() {
let config = rewrite_config();
let mut body = json!({
"result": {
"data": "unchanged",
"_meta": { "openai/widgetDomain": "should-stay.com" }
}
});
let _ = rewrite_response("notifications/message", &mut body, &config);
assert_eq!(
body["result"]["_meta"]["openai/widgetDomain"]
.as_str()
.unwrap(),
"should-stay.com"
);
assert_eq!(body["result"]["data"].as_str().unwrap(), "unchanged");
}
#[test]
fn rewrite_response__replace_mode_ignores_upstream() {
let mut config = rewrite_config();
config.csp.resource_domains = DirectivePolicy {
domains: vec!["https://allowed.example.com".into()],
mode: Mode::Replace,
};
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://allowed.example.com".into()],
mode: Mode::Replace,
};
let mut body = json!({
"result": {
"tools": [{
"name": "test",
"_meta": {
"openai/widgetCSP": {
"resource_domains": ["https://cdn.external.com", "https://api.external.com"],
"connect_domains": ["https://api.external.com", "http://localhost:9000"]
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let resources =
as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["resource_domains"]);
assert_eq!(
resources,
vec![
"https://abc.tunnel.example.com",
"https://allowed.example.com"
]
);
let connect =
as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
assert_eq!(
connect,
vec![
"https://abc.tunnel.example.com",
"https://allowed.example.com"
]
);
}
#[test]
fn rewrite_response__widget_scope_matches_resource_uri() {
let mut config = rewrite_config();
config.csp.widgets.push(WidgetScoped {
match_pattern: "ui://widget/payment*".into(),
connect_domains: vec!["https://api.stripe.com".into()],
connect_domains_mode: Mode::Extend,
..Default::default()
});
let mut body = json!({
"result": {
"resources": [
{
"uri": "ui://widget/payment-form",
"_meta": {
"openai/widgetCSP": { "connect_domains": [] }
}
},
{
"uri": "ui://widget/search",
"_meta": {
"openai/widgetCSP": { "connect_domains": [] }
}
}
]
}
});
let _ = rewrite_response("resources/list", &mut body, &config);
let payment_connect = as_strs(
&body["result"]["resources"][0]["_meta"]["openai/widgetCSP"]["connect_domains"],
);
assert!(payment_connect.contains(&"https://api.stripe.com"));
let search_connect = as_strs(
&body["result"]["resources"][1]["_meta"]["openai/widgetCSP"]["connect_domains"],
);
assert!(!search_connect.contains(&"https://api.stripe.com"));
}
#[test]
fn rewrite_response__widget_replace_mode_wipes_upstream() {
let mut config = rewrite_config();
config.csp.widgets.push(WidgetScoped {
match_pattern: "ui://widget/*".into(),
connect_domains: vec!["https://api.stripe.com".into()],
connect_domains_mode: Mode::Replace,
..Default::default()
});
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/payment",
"_meta": {
"openai/widgetCSP": {
"connect_domains": [
"https://api.external.com",
"https://another.external.com"
]
}
}
}]
}
});
let _ = rewrite_response("resources/read", &mut body, &config);
let connect =
as_strs(&body["result"]["contents"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
assert_eq!(
connect,
vec!["https://abc.tunnel.example.com", "https://api.stripe.com"]
);
}
#[test]
fn rewrite_response__widget_uri_inferred_from_tool_meta() {
let mut config = rewrite_config();
config.csp.widgets.push(WidgetScoped {
match_pattern: "ui://widget/payment*".into(),
connect_domains: vec!["https://api.stripe.com".into()],
connect_domains_mode: Mode::Extend,
..Default::default()
});
let mut body = json!({
"result": {
"tools": [{
"name": "take_payment",
"_meta": {
"ui": { "resourceUri": "ui://widget/payment-form" },
"openai/widgetCSP": { "connect_domains": [] }
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let connect =
as_strs(&body["result"]["tools"][0]["_meta"]["openai/widgetCSP"]["connect_domains"]);
assert!(connect.contains(&"https://api.stripe.com"));
}
#[test]
fn rewrite_response__spec_only_upstream_also_emits_openai_shape() {
let config = rewrite_config();
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/search",
"mimeType": "text/html",
"_meta": {
"ui": {
"csp": {
"connectDomains": ["https://api.external.com"],
"resourceDomains": ["https://cdn.external.com"]
}
}
}
}]
}
});
let _ = rewrite_response("resources/read", &mut body, &config);
let meta = &body["result"]["contents"][0]["_meta"];
let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa_connect, spec_connect);
assert!(oa_connect.contains(&"https://api.external.com"));
assert!(oa_connect.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__openai_only_upstream_also_emits_spec_shape() {
let config = rewrite_config();
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/search",
"mimeType": "text/html",
"_meta": {
"openai/widgetCSP": {
"connect_domains": ["https://api.external.com"],
"resource_domains": ["https://cdn.external.com"]
}
}
}]
}
});
let _ = rewrite_response("resources/read", &mut body, &config);
let meta = &body["result"]["contents"][0]["_meta"];
let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa_connect, spec_connect);
assert!(spec_connect.contains(&"https://api.external.com"));
assert!(spec_connect.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__declared_config_synthesizes_both_shapes_from_empty() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.declared.com".into()],
mode: Mode::Extend,
};
let mut body = json!({
"result": {
"resources": [{
"uri": "ui://widget/search",
"_meta": {
"openai/widgetDomain": "old.domain.com"
}
}]
}
});
let _ = rewrite_response("resources/list", &mut body, &config);
let meta = &body["result"]["resources"][0]["_meta"];
let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa, spec);
assert!(oa.contains(&"https://api.declared.com"));
assert!(oa.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__upstream_declarations_unioned_across_shapes() {
let config = rewrite_config();
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/search",
"mimeType": "text/html",
"_meta": {
"openai/widgetCSP": {
"connect_domains": ["https://api.only-openai.com"]
},
"ui": {
"csp": {
"connectDomains": ["https://api.only-spec.com"]
}
}
}
}]
}
});
let _ = rewrite_response("resources/read", &mut body, &config);
let meta = &body["result"]["contents"][0]["_meta"];
let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa, spec);
assert!(oa.contains(&"https://api.only-openai.com"));
assert!(oa.contains(&"https://api.only-spec.com"));
}
#[test]
fn rewrite_response__non_widget_meta_is_not_polluted() {
let config = rewrite_config();
let mut body = json!({
"result": {
"content": [{"type": "text", "text": "plain result"}],
"_meta": { "requestId": "abc-123" }
}
});
let _ = rewrite_response("tools/call", &mut body, &config);
let meta = &body["result"]["_meta"];
assert!(meta.get("openai/widgetCSP").is_none());
assert!(meta.get("ui").is_none());
assert_eq!(meta["requestId"].as_str().unwrap(), "abc-123");
}
#[test]
fn rewrite_response__all_three_directives_synthesized() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.example.com".into()],
mode: Mode::Extend,
};
config.csp.resource_domains = DirectivePolicy {
domains: vec!["https://cdn.example.com".into()],
mode: Mode::Extend,
};
let mut body = json!({
"result": {
"resources": [{
"uri": "ui://widget/search",
"_meta": { "openai/widgetDomain": "x" }
}]
}
});
let _ = rewrite_response("resources/list", &mut body, &config);
let meta = &body["result"]["resources"][0]["_meta"];
for shape in ["openai/widgetCSP"] {
assert!(meta[shape]["connect_domains"].is_array());
assert!(meta[shape]["resource_domains"].is_array());
assert!(meta[shape]["frame_domains"].is_array());
}
assert!(meta["ui"]["csp"]["connectDomains"].is_array());
assert!(meta["ui"]["csp"]["resourceDomains"].is_array());
assert!(meta["ui"]["csp"]["frameDomains"].is_array());
}
#[test]
fn rewrite_response__frame_domains_default_replace_drops_upstream() {
let config = rewrite_config();
let mut body = json!({
"result": {
"tools": [{
"name": "test",
"_meta": {
"ui": {
"csp": {
"frameDomains": ["https://embed.external.com"]
}
}
}
}]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let frames = as_strs(&body["result"]["tools"][0]["_meta"]["ui"]["csp"]["frameDomains"]);
assert_eq!(frames, vec!["https://abc.tunnel.example.com"]);
}
#[test]
fn rewrite_response__end_to_end_mcp_schema() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.myshop.com".into()],
mode: Mode::Extend,
};
config.csp.resource_domains = DirectivePolicy {
domains: vec!["https://cdn.myshop.com".into()],
mode: Mode::Extend,
};
config.csp.widgets.push(WidgetScoped {
match_pattern: "ui://widget/payment*".into(),
connect_domains: vec!["https://api.stripe.com".into()],
connect_domains_mode: Mode::Extend,
resource_domains: vec!["https://js.stripe.com".into()],
resource_domains_mode: Mode::Extend,
..Default::default()
});
let mut body = json!({
"jsonrpc": "2.0",
"id": 42,
"result": {
"tools": [
{
"name": "search_products",
"description": "Search the product catalog",
"inputSchema": { "type": "object" },
"_meta": {
"openai/widgetDomain": "old.shop.com",
"openai/outputTemplate": "ui://widget/search",
"openai/widgetCSP": {
"connect_domains": ["http://localhost:9000"],
"resource_domains": ["http://localhost:4444"]
}
}
},
{
"name": "take_payment",
"description": "Charge a card",
"inputSchema": { "type": "object" },
"_meta": {
"ui": {
"resourceUri": "ui://widget/payment-form",
"csp": {
"connectDomains": ["https://api.myshop.com"]
}
}
}
},
{
"name": "get_order_status",
"description": "Look up an order",
"inputSchema": { "type": "object" }
}
]
}
});
let _ = rewrite_response("tools/list", &mut body, &config);
let tools = body["result"]["tools"].as_array().unwrap();
let search_meta = &tools[0]["_meta"];
assert_eq!(
search_meta["openai/widgetDomain"].as_str().unwrap(),
"abc.tunnel.example.com"
);
let search_oa_connect = as_strs(&search_meta["openai/widgetCSP"]["connect_domains"]);
let search_spec_connect = as_strs(&search_meta["ui"]["csp"]["connectDomains"]);
assert_eq!(search_oa_connect, search_spec_connect);
assert_eq!(
search_oa_connect,
vec!["https://abc.tunnel.example.com", "https://api.myshop.com"]
);
assert!(!search_oa_connect.contains(&"https://api.stripe.com"));
let search_oa_frame = as_strs(&search_meta["openai/widgetCSP"]["frame_domains"]);
assert_eq!(search_oa_frame, vec!["https://abc.tunnel.example.com"]);
let payment_meta = &tools[1]["_meta"];
let payment_oa_connect = as_strs(&payment_meta["openai/widgetCSP"]["connect_domains"]);
let payment_spec_connect = as_strs(&payment_meta["ui"]["csp"]["connectDomains"]);
assert_eq!(payment_oa_connect, payment_spec_connect);
assert_eq!(
payment_oa_connect,
vec![
"https://abc.tunnel.example.com",
"https://api.myshop.com",
"https://api.stripe.com",
]
);
let payment_oa_resource = as_strs(&payment_meta["openai/widgetCSP"]["resource_domains"]);
assert_eq!(
payment_oa_resource,
vec![
"https://abc.tunnel.example.com",
"https://cdn.myshop.com",
"https://js.stripe.com",
]
);
let plain = &tools[2];
assert!(plain.get("_meta").is_none());
}
#[test]
fn rewrite_response__tools_call_underscore_meta_is_rewritten() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://assets.usestudykit.com".into()],
mode: Mode::Replace,
};
config.csp.resource_domains = DirectivePolicy {
domains: vec!["https://assets.usestudykit.com".into()],
mode: Mode::Replace,
};
let mut body = json!({
"result": {
"_meta": {
"openai/outputTemplate": "ui://widget/vocab_review.html",
"openai/widgetDomain": "assets.usestudykit.com/src",
"openai/widgetCSP": {
"connect_domains": [
"http://localhost:9002",
"https://api.dictionaryapi.dev"
],
"resource_domains": [
"http://localhost:9002",
"https://api.dictionaryapi.dev"
]
},
"ui": {
"csp": {
"connectDomains": ["https://api.dictionaryapi.dev"],
"resourceDomains": ["https://api.dictionaryapi.dev"]
},
"resourceUri": "ui://widget/vocab_review.html"
}
},
"content": [{"type": "text", "text": "payload"}],
"structuredContent": {"data": {"items": []}}
}
});
let _ = rewrite_response("tools/call", &mut body, &config);
let meta = &body["result"]["_meta"];
assert_eq!(
meta["openai/widgetDomain"].as_str().unwrap(),
"abc.tunnel.example.com"
);
let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa_connect, spec_connect);
assert_eq!(
oa_connect,
vec![
"https://abc.tunnel.example.com",
"https://assets.usestudykit.com"
]
);
let oa_resource = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
assert_eq!(
oa_resource,
vec![
"https://abc.tunnel.example.com",
"https://assets.usestudykit.com"
]
);
assert_eq!(
body["result"]["content"][0]["text"].as_str().unwrap(),
"payload"
);
}
#[test]
fn rewrite_response__resources_read_underscore_meta_is_rewritten() {
let config = rewrite_config();
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/question",
"mimeType": "text/html",
"text": "<html/>",
"_meta": {
"openai/widgetDomain": "old.domain.com"
}
}]
}
});
let _ = rewrite_response("resources/read", &mut body, &config);
assert_eq!(
body["result"]["contents"][0]["_meta"]["openai/widgetDomain"]
.as_str()
.unwrap(),
"abc.tunnel.example.com"
);
}
#[test]
fn rewrite_response__legacy_meta_key_is_ignored() {
let config = rewrite_config();
let mut body = json!({
"result": {
"_meta": {"openai/widgetDomain": "real.domain.com"},
"meta": {"openai/widgetDomain": "should-stay.com"}
}
});
let _ = rewrite_response("tools/call", &mut body, &config);
assert_eq!(
body["result"]["_meta"]["openai/widgetDomain"]
.as_str()
.unwrap(),
"abc.tunnel.example.com"
);
assert_eq!(
body["result"]["meta"]["openai/widgetDomain"]
.as_str()
.unwrap(),
"should-stay.com"
);
}
#[test]
fn rewrite_response__resources_list_synthesizes_meta_when_upstream_omits() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.declared.com".into()],
mode: Mode::Replace,
};
config.csp.resource_domains = DirectivePolicy {
domains: vec!["https://cdn.declared.com".into()],
mode: Mode::Extend,
};
let mut body = json!({
"result": {
"resources": [{
"uri": "ui://widget/search",
"name": "Search Widget"
}]
}
});
let mutated = rewrite_response("resources/list", &mut body, &config);
assert!(mutated);
let meta = &body["result"]["resources"][0]["_meta"];
let oa_connect = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec_connect = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa_connect, spec_connect);
assert_eq!(
oa_connect,
vec!["https://abc.tunnel.example.com", "https://api.declared.com"]
);
let oa_resource = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
assert!(oa_resource.contains(&"https://abc.tunnel.example.com"));
assert!(oa_resource.contains(&"https://cdn.declared.com"));
}
#[test]
fn rewrite_response__resources_read_synthesizes_meta_when_upstream_omits() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.declared.com".into()],
mode: Mode::Replace,
};
let mut body = json!({
"result": {
"contents": [{
"uri": "ui://widget/question",
"mimeType": "text/html",
"text": "<html><body>Hello</body></html>"
}]
}
});
let mutated = rewrite_response("resources/read", &mut body, &config);
assert!(mutated);
assert_eq!(
body["result"]["contents"][0]["text"].as_str().unwrap(),
"<html><body>Hello</body></html>"
);
let meta = &body["result"]["contents"][0]["_meta"];
let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
let spec = as_strs(&meta["ui"]["csp"]["connectDomains"]);
assert_eq!(oa, spec);
assert_eq!(
oa,
vec!["https://abc.tunnel.example.com", "https://api.declared.com"]
);
}
#[test]
fn rewrite_response__resources_list_injects_into_empty_meta() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.declared.com".into()],
mode: Mode::Extend,
};
let mut body = json!({
"result": {
"resources": [{
"uri": "ui://widget/search",
"_meta": {}
}]
}
});
let _ = rewrite_response("resources/list", &mut body, &config);
let meta = &body["result"]["resources"][0]["_meta"];
let oa = as_strs(&meta["openai/widgetCSP"]["connect_domains"]);
assert!(oa.contains(&"https://api.declared.com"));
assert!(oa.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__resources_templates_list_synthesizes_meta() {
let mut config = rewrite_config();
config.csp.resource_domains = DirectivePolicy {
domains: vec!["https://cdn.declared.com".into()],
mode: Mode::Extend,
};
let mut body = json!({
"result": {
"resourceTemplates": [{
"uriTemplate": "ui://widget/{name}.html",
"name": "Widget Template"
}]
}
});
let _ = rewrite_response("resources/templates/list", &mut body, &config);
let meta = &body["result"]["resourceTemplates"][0]["_meta"];
let oa = as_strs(&meta["openai/widgetCSP"]["resource_domains"]);
assert!(oa.contains(&"https://cdn.declared.com"));
assert!(oa.contains(&"https://abc.tunnel.example.com"));
}
#[test]
fn rewrite_response__tools_call_no_meta_is_not_synthesized() {
let mut config = rewrite_config();
config.csp.connect_domains = DirectivePolicy {
domains: vec!["https://api.declared.com".into()],
mode: Mode::Replace,
};
let mut body = json!({
"result": {
"content": [{"type": "text", "text": "London 14C"}],
"structuredContent": {"city": "London", "temp": 14}
}
});
let _ = rewrite_response("tools/call", &mut body, &config);
assert!(body["result"].get("_meta").is_none());
assert_eq!(
body["result"]["content"][0]["text"].as_str().unwrap(),
"London 14C"
);
}
#[test]
fn rewrite_response__resources_list_skips_when_no_uri_and_no_meta() {
let config = rewrite_config();
let mut body = json!({
"result": {
"resources": [{
"name": "malformed"
}]
}
});
let _ = rewrite_response("resources/list", &mut body, &config);
assert!(body["result"]["resources"][0].get("_meta").is_none());
}
}