use std::collections::HashMap;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use super::agent::AgentExtension;
use super::completion::CompletionExtension;
use super::delegation::DelegationExtension;
use super::framework::FrameworkExtension;
use super::guarded::{Guarded, WriteToken};
use super::http::HttpExtension;
use super::llm::LLMExtension;
use super::mcp::MCPExtension;
use super::meta::MetaExtension;
use super::provenance::ProvenanceExtension;
use super::raw_credentials::RawCredentialsExtension;
use super::request::RequestExtension;
use super::security::SecurityExtension;
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct Extensions {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request: Option<Arc<RequestExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub agent: Option<Arc<AgentExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub http: Option<Arc<HttpExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub security: Option<Arc<SecurityExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub delegation: Option<Arc<DelegationExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub raw_credentials: Option<Arc<RawCredentialsExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mcp: Option<Arc<MCPExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub completion: Option<Arc<CompletionExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provenance: Option<Arc<ProvenanceExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub llm: Option<Arc<LLMExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub framework: Option<Arc<FrameworkExtension>>,
#[serde(default)]
pub meta: Option<Arc<MetaExtension>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub custom: Option<Arc<HashMap<String, serde_json::Value>>>,
#[serde(skip)]
pub http_write_token: Option<WriteToken>,
#[serde(skip)]
pub labels_write_token: Option<WriteToken>,
#[serde(skip)]
pub delegation_write_token: Option<WriteToken>,
}
impl Clone for Extensions {
fn clone(&self) -> Self {
Self {
request: self.request.clone(),
agent: self.agent.clone(),
http: self.http.clone(),
security: self.security.clone(),
delegation: self.delegation.clone(),
raw_credentials: self.raw_credentials.clone(),
mcp: self.mcp.clone(),
completion: self.completion.clone(),
provenance: self.provenance.clone(),
llm: self.llm.clone(),
framework: self.framework.clone(),
meta: self.meta.clone(),
custom: self.custom.clone(),
http_write_token: None,
labels_write_token: None,
delegation_write_token: None,
}
}
}
impl Extensions {
pub fn cow_copy(&self) -> OwnedExtensions {
OwnedExtensions {
request: self.request.clone(),
agent: self.agent.clone(),
mcp: self.mcp.clone(),
completion: self.completion.clone(),
provenance: self.provenance.clone(),
llm: self.llm.clone(),
framework: self.framework.clone(),
meta: self.meta.clone(),
raw_credentials: self.raw_credentials.clone(),
http: self.http.as_ref().map(|arc| Guarded::new((**arc).clone())),
security: self.security.as_ref().map(|arc| (**arc).clone()),
delegation: self.delegation.as_ref().map(|arc| (**arc).clone()),
custom: self.custom.as_ref().map(|arc| (**arc).clone()),
http_write_token: if self.http_write_token.is_some() {
Some(WriteToken::new())
} else {
None
},
labels_write_token: if self.labels_write_token.is_some() {
Some(WriteToken::new())
} else {
None
},
delegation_write_token: if self.delegation_write_token.is_some() {
Some(WriteToken::new())
} else {
None
},
}
}
pub fn validate_immutable(&self, modified: &OwnedExtensions) -> bool {
fn ptr_eq_opt<T>(a: &Option<Arc<T>>, b: &Option<Arc<T>>) -> bool {
match (a, b) {
(Some(a), Some(b)) => Arc::ptr_eq(a, b),
(None, None) => true,
(_, None) => true, (None, Some(_)) => false, }
}
ptr_eq_opt(&self.request, &modified.request)
&& ptr_eq_opt(&self.agent, &modified.agent)
&& ptr_eq_opt(&self.mcp, &modified.mcp)
&& ptr_eq_opt(&self.completion, &modified.completion)
&& ptr_eq_opt(&self.provenance, &modified.provenance)
&& ptr_eq_opt(&self.llm, &modified.llm)
&& ptr_eq_opt(&self.framework, &modified.framework)
&& ptr_eq_opt(&self.meta, &modified.meta)
}
pub fn merge_owned(&mut self, owned: OwnedExtensions) {
self.http = owned.http.map(|g| Arc::new(g.into_inner()));
self.security = owned.security.map(Arc::new);
self.delegation = owned.delegation.map(Arc::new);
self.custom = owned.custom.map(Arc::new);
if owned.raw_credentials.is_some() {
self.raw_credentials = owned.raw_credentials;
}
}
}
#[derive(Debug)]
pub struct OwnedExtensions {
pub request: Option<Arc<RequestExtension>>,
pub agent: Option<Arc<AgentExtension>>,
pub mcp: Option<Arc<MCPExtension>>,
pub completion: Option<Arc<CompletionExtension>>,
pub provenance: Option<Arc<ProvenanceExtension>>,
pub llm: Option<Arc<LLMExtension>>,
pub framework: Option<Arc<FrameworkExtension>>,
pub meta: Option<Arc<MetaExtension>>,
pub raw_credentials: Option<Arc<RawCredentialsExtension>>,
pub http: Option<Guarded<HttpExtension>>,
pub security: Option<SecurityExtension>,
pub delegation: Option<DelegationExtension>,
pub custom: Option<HashMap<String, serde_json::Value>>,
pub http_write_token: Option<WriteToken>,
pub labels_write_token: Option<WriteToken>,
pub delegation_write_token: Option<WriteToken>,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::extensions::{
DelegationExtension, HttpExtension, RequestExtension, SecurityExtension,
};
fn make_extensions() -> Extensions {
let mut security = SecurityExtension::default();
security.add_label("PII");
let mut http = HttpExtension::default();
http.set_header("Authorization", "Bearer token");
Extensions {
request: Some(Arc::new(RequestExtension {
request_id: Some("req-001".into()),
..Default::default()
})),
security: Some(Arc::new(security)),
http: Some(Arc::new(http)),
delegation: Some(Arc::new(DelegationExtension::default())),
meta: Some(Arc::new(MetaExtension {
entity_type: Some("tool".into()),
..Default::default()
})),
..Default::default()
}
}
#[test]
fn test_cow_copy_shares_immutable_arcs() {
let ext = make_extensions();
let cow = ext.cow_copy();
assert!(Arc::ptr_eq(
ext.request.as_ref().unwrap(),
cow.request.as_ref().unwrap()
));
assert!(Arc::ptr_eq(
ext.meta.as_ref().unwrap(),
cow.meta.as_ref().unwrap()
));
}
#[test]
fn test_cow_copy_deep_clones_mutable_slots() {
let ext = make_extensions();
let cow = ext.cow_copy();
assert!(cow.security.is_some());
assert!(cow.http.is_some());
assert!(cow.delegation.is_some());
cow.security.as_ref().unwrap().has_label("PII");
}
#[test]
fn test_cow_copy_propagates_write_tokens() {
let mut ext = make_extensions();
let cow_no_tokens = ext.cow_copy();
assert!(cow_no_tokens.http_write_token.is_none());
assert!(cow_no_tokens.labels_write_token.is_none());
assert!(cow_no_tokens.delegation_write_token.is_none());
ext.http_write_token = Some(WriteToken::new());
ext.labels_write_token = Some(WriteToken::new());
let cow_with_tokens = ext.cow_copy();
assert!(cow_with_tokens.http_write_token.is_some());
assert!(cow_with_tokens.labels_write_token.is_some());
assert!(cow_with_tokens.delegation_write_token.is_none()); }
#[test]
fn test_cow_copy_write_token_enables_guarded_write() {
let mut ext = make_extensions();
ext.http_write_token = Some(WriteToken::new());
let mut cow = ext.cow_copy();
assert_eq!(
cow.http
.as_ref()
.unwrap()
.read()
.get_header("Authorization"),
Some("Bearer token")
);
let token = cow.http_write_token.as_ref().unwrap();
cow.http
.as_mut()
.unwrap()
.write(token)
.set_header("X-Custom", "value");
assert_eq!(
cow.http.as_ref().unwrap().read().get_header("X-Custom"),
Some("value")
);
assert!(ext.http.as_ref().unwrap().get_header("X-Custom").is_none());
}
#[test]
fn test_cow_copy_monotonic_label_insert() {
let mut ext = make_extensions();
ext.labels_write_token = Some(WriteToken::new());
let mut cow = ext.cow_copy();
cow.security.as_mut().unwrap().add_label("HIPAA");
assert!(cow.security.as_ref().unwrap().has_label("HIPAA"));
assert!(!ext.security.as_ref().unwrap().has_label("HIPAA"));
}
#[test]
fn test_validate_immutable_passes_for_cow() {
let ext = make_extensions();
let cow = ext.cow_copy();
assert!(ext.validate_immutable(&cow));
}
#[test]
fn test_validate_immutable_fails_when_tampered() {
let ext = make_extensions();
let mut cow = ext.cow_copy();
cow.request = Some(Arc::new(RequestExtension {
request_id: Some("TAMPERED".into()),
..Default::default()
}));
assert!(!ext.validate_immutable(&cow));
}
#[test]
fn test_validate_immutable_both_none_passes() {
let ext = Extensions::default();
let cow = ext.cow_copy();
assert!(ext.validate_immutable(&cow));
}
#[test]
fn test_clone_drops_write_tokens() {
let mut ext = make_extensions();
ext.http_write_token = Some(WriteToken::new());
ext.labels_write_token = Some(WriteToken::new());
ext.delegation_write_token = Some(WriteToken::new());
let cloned = ext.clone();
assert!(cloned.http_write_token.is_none());
assert!(cloned.labels_write_token.is_none());
assert!(cloned.delegation_write_token.is_none());
let cow = ext.cow_copy();
assert!(cow.http_write_token.is_some());
assert!(cow.labels_write_token.is_some());
assert!(cow.delegation_write_token.is_some());
}
#[test]
fn test_cow_copy_modify_multiple_fields() {
use crate::extensions::delegation::DelegationHop;
use crate::extensions::DelegationExtension;
let mut security = SecurityExtension::default();
security.add_label("PII");
let mut http = HttpExtension::default();
http.set_header("Authorization", "Bearer token");
let mut ext = Extensions {
security: Some(Arc::new(security)),
http: Some(Arc::new(http)),
delegation: Some(Arc::new(DelegationExtension::default())),
custom: Some(Arc::new(
[("existing".to_string(), serde_json::json!("value"))].into(),
)),
meta: Some(Arc::new(MetaExtension {
entity_type: Some("tool".into()),
..Default::default()
})),
..Default::default()
};
ext.http_write_token = Some(WriteToken::new());
ext.labels_write_token = Some(WriteToken::new());
ext.delegation_write_token = Some(WriteToken::new());
let mut cow = ext.cow_copy();
cow.security.as_mut().unwrap().add_label("CHECKED");
cow.security.as_mut().unwrap().add_label("COMPLIANT");
let token = cow.http_write_token.as_ref().unwrap();
cow.http
.as_mut()
.unwrap()
.write(token)
.set_header("X-Checked", "true");
cow.http
.as_mut()
.unwrap()
.write(token)
.set_header("X-Policy", "v2");
cow.delegation.as_mut().unwrap().append_hop(DelegationHop {
subject_id: "service-a".into(),
scopes_granted: vec!["read_hr".into()],
..Default::default()
});
cow.custom
.as_mut()
.unwrap()
.insert("audit.timestamp".into(), serde_json::json!("2026-04-29"));
let sec = cow.security.as_ref().unwrap();
assert!(sec.has_label("PII")); assert!(sec.has_label("CHECKED")); assert!(sec.has_label("COMPLIANT"));
let http = cow.http.as_ref().unwrap().read();
assert_eq!(http.get_header("Authorization"), Some("Bearer token")); assert_eq!(http.get_header("X-Checked"), Some("true")); assert_eq!(http.get_header("X-Policy"), Some("v2"));
assert_eq!(cow.delegation.as_ref().unwrap().chain.len(), 1);
assert_eq!(
cow.delegation.as_ref().unwrap().chain[0].subject_id,
"service-a"
);
assert_eq!(
cow.custom.as_ref().unwrap().get("existing").unwrap(),
"value"
);
assert_eq!(
cow.custom.as_ref().unwrap().get("audit.timestamp").unwrap(),
"2026-04-29"
);
assert!(!ext.security.as_ref().unwrap().has_label("CHECKED"));
assert!(ext.http.as_ref().unwrap().get_header("X-Checked").is_none());
assert!(ext.delegation.as_ref().unwrap().chain.is_empty());
assert!(!ext.custom.as_ref().unwrap().contains_key("audit.timestamp"));
assert!(ext.validate_immutable(&cow));
}
#[test]
fn test_validate_immutable_passes_when_slot_filtered_out() {
let ext = make_extensions();
let mut cow = ext.cow_copy();
cow.agent = None;
assert!(ext.validate_immutable(&cow));
}
#[test]
fn test_validate_immutable_fails_when_slot_fabricated() {
let ext = Extensions::default(); let mut cow = ext.cow_copy();
cow.agent = Some(Arc::new(AgentExtension {
agent_id: Some("fabricated".into()),
..Default::default()
}));
assert!(!ext.validate_immutable(&cow));
}
#[test]
fn test_validate_immutable_passes_multiple_slots_filtered() {
let ext = make_extensions();
let mut cow = ext.cow_copy();
cow.agent = None;
cow.mcp = None;
cow.completion = None;
cow.framework = None;
assert!(ext.validate_immutable(&cow));
}
#[test]
fn test_merge_owned_preserves_http_response_headers() {
let mut http = HttpExtension::default();
http.set_request_header("Authorization", "Bearer tok");
let mut ext = Extensions {
http: Some(Arc::new(http)),
..Default::default()
};
ext.http_write_token = Some(WriteToken::new());
let mut cow = ext.cow_copy();
let token = cow.http_write_token.as_ref().unwrap();
let h = cow.http.as_mut().unwrap().write(token);
h.set_response_header("X-Tool-Name", "get_compensation");
h.set_response_header("X-Status", "success");
ext.merge_owned(cow);
let merged_http = ext.http.as_ref().unwrap();
assert_eq!(
merged_http.get_response_header("X-Tool-Name"),
Some("get_compensation")
);
assert_eq!(merged_http.get_response_header("X-Status"), Some("success"));
assert_eq!(
merged_http.get_request_header("Authorization"),
Some("Bearer tok")
);
}
#[test]
fn test_merge_owned_with_filtered_security() {
let mut security = SecurityExtension::default();
security.add_label("PII");
security.add_label("HR");
let ext = Extensions {
security: Some(Arc::new(security)),
..Default::default()
};
let mut cow = ext.cow_copy();
cow.security.as_mut().unwrap().labels = crate::extensions::MonotonicSet::new();
let mut ext_mut = ext.clone();
ext_mut.merge_owned(cow);
let merged_sec = ext_mut.security.as_ref().unwrap();
assert!(!merged_sec.has_label("PII")); }
#[test]
fn test_merge_owned_none_http_preserves_pipeline() {
let mut http = HttpExtension::default();
http.set_request_header("X-Original", "value");
let mut ext = Extensions {
http: Some(Arc::new(http)),
..Default::default()
};
let mut cow = ext.cow_copy();
cow.http = None;
ext.merge_owned(cow);
assert!(ext.http.is_none());
}
#[test]
fn test_read_only_plugin_zero_cost() {
let ext = make_extensions();
let has_pii = ext
.security
.as_ref()
.map(|s| s.has_label("PII"))
.unwrap_or(false);
assert!(has_pii);
let auth = ext
.http
.as_ref()
.and_then(|h| h.get_header("Authorization"));
assert_eq!(auth, Some("Bearer token"));
let entity = ext.meta.as_ref().and_then(|m| m.entity_type.as_deref());
assert_eq!(entity, Some("tool"));
}
}