use std::sync::Arc;
use crate::trigger_engine::notification_hook::{
HookError, HookState, NotificationHook, NotificationHookStatus, TriggerSink,
};
use crate::trigger_engine::types::{
CredentialScope, PayloadVisibility, ReplacementPolicy, SourceKind, Trigger, TriggerAuthority,
TriggerSource,
};
use async_trait::async_trait;
use chrono::Utc;
use parking_lot::Mutex;
use sha2::{Digest, Sha256};
use theway_mcp::client::McpServerNotification;
use tokio::sync::mpsc::UnboundedReceiver;
use uuid::Uuid;
pub struct McpNotificationHook {
label: String,
server_name: String,
rx: Mutex<Option<UnboundedReceiver<McpServerNotification>>>,
status: Arc<Mutex<NotificationHookStatus>>,
}
impl McpNotificationHook {
pub fn new(
server_name: impl Into<String>,
rx: UnboundedReceiver<McpServerNotification>,
) -> Self {
let server_name = server_name.into();
let label = format!("mcp:{server_name}");
let mut status = NotificationHookStatus::pending();
status.subscription_labels = vec![label.clone()];
Self {
label,
server_name,
rx: Mutex::new(Some(rx)),
status: Arc::new(Mutex::new(status)),
}
}
#[cfg(test)]
fn debug_status_handle(&self) -> Arc<Mutex<NotificationHookStatus>> {
self.status.clone()
}
}
#[async_trait]
impl NotificationHook for McpNotificationHook {
fn label(&self) -> &str {
&self.label
}
async fn run(&self, sink: TriggerSink) -> Result<(), HookError> {
let mut rx = self.rx.lock().take().ok_or_else(|| {
HookError::Other(format!(
"{} hook already ran; receiver consumed",
self.label
))
})?;
self.status.lock().state = HookState::Connected;
while let Some(notification) = rx.recv().await {
let trigger = match map_notification(&self.server_name, ¬ification) {
Some(t) => t,
None => {
let mut st = self.status.lock();
st.dropped_count = st.dropped_count.saturating_add(1);
st.last_error = Some(format!(
"dropped custom notification {:?}: missing `_meta.theway_dedup_key`",
notification.method
));
continue;
}
};
if sink.send(trigger).is_err() {
self.status.lock().state = HookState::Disconnected {
reason: "sink closed".into(),
};
return Err(HookError::SinkClosed);
}
let mut st = self.status.lock();
st.last_event_at = Some(Utc::now());
st.last_error = None;
}
self.status.lock().state = HookState::Disconnected {
reason: "mcp transport closed".into(),
};
Ok(())
}
fn status(&self) -> NotificationHookStatus {
self.status.lock().clone()
}
}
fn map_notification(server_name: &str, n: &McpServerNotification) -> Option<Trigger> {
let (idempotency_key, replacement_policy) = idempotency_for(server_name, &n.method, &n.params)?;
let payload_summary = render_summary(&n.method, &n.params);
Some(Trigger {
source: TriggerSource::Mcp {
server_name: server_name.to_string(),
method: n.method.clone(),
},
source_kind: SourceKind::Mcp,
source_label: format!("mcp:{server_name}"),
event_label: n.method.clone(),
payload_visibility: PayloadVisibility::Local,
payload_summary,
payload: None,
idempotency_key,
replacement_policy,
trace_id: Uuid::new_v4().to_string(),
authority: TriggerAuthority {
principal_id: format!("mcp:{server_name}"),
principal_label: server_name.to_string(),
credential_scope: CredentialScope::User,
allowed_source_actions: Vec::new(),
expires_at: None,
},
received_at: Utc::now(),
})
}
pub(crate) fn safe_display(value: &str, cap: usize) -> String {
let redacted = redact_notification_text(value).replace('\n', " ");
truncate_chars(&redacted, cap)
}
fn redact_notification_text(value: &str) -> String {
value
.split_whitespace()
.map(|part| {
let lower = part.to_ascii_lowercase();
if lower.starts_with("hub_agent_")
|| lower.starts_with("hub_hs_")
|| lower.starts_with("hub_ep_")
|| lower.starts_with("sk-")
|| lower.contains("bearer")
|| lower.contains("token")
{
"[redacted]"
} else {
part
}
})
.collect::<Vec<_>>()
.join(" ")
}
fn truncate_chars(value: &str, cap: usize) -> String {
if value.chars().count() <= cap {
return value.to_string();
}
let mut out = value
.chars()
.take(cap.saturating_sub(1))
.collect::<String>();
out.push('…');
out
}
fn idempotency_for(
server_name: &str,
method: &str,
params: &serde_json::Value,
) -> Option<(String, ReplacementPolicy)> {
let prefix = format!("mcp:{server_name}:");
match method {
"notifications/tools/listChanged" => {
Some((format!("{prefix}tools"), ReplacementPolicy::LatestReplaces))
}
"notifications/resources/listChanged" => Some((
format!("{prefix}resources"),
ReplacementPolicy::LatestReplaces,
)),
"notifications/prompts/listChanged" => Some((
format!("{prefix}prompts"),
ReplacementPolicy::LatestReplaces,
)),
"notifications/resources/updated" => {
let uri = params
.get("uri")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
Some((
format!("{prefix}resources:{}", safe_idempotency_segment(uri)),
ReplacementPolicy::LatestReplaces,
))
}
_ => {
extract_dedup_key(params).map(|k| {
(
format!("{prefix}custom:{}", safe_idempotency_segment(&k)),
ReplacementPolicy::Drop,
)
})
}
}
}
fn extract_dedup_key(params: &serde_json::Value) -> Option<String> {
params
.get("_meta")
.and_then(|m| m.get("theway_dedup_key"))
.and_then(|v| v.as_str())
.map(|s| s.to_string())
}
pub(crate) fn safe_idempotency_segment(value: &str) -> String {
let redacted = redact_notification_text(value);
let has_sensitive_text = redacted != value;
let is_unbounded = value.chars().count() > 200;
let has_control_chars = value.chars().any(|ch| ch.is_control());
if has_sensitive_text || is_unbounded || has_control_chars {
let digest = Sha256::digest(value.as_bytes());
return format!("hash:{}", hex::encode(&digest[..6]));
}
value.to_string()
}
fn render_summary(method: &str, params: &serde_json::Value) -> Option<String> {
match method {
"notifications/resources/updated" => {
if let Some(uri) = params.get("uri").and_then(|v| v.as_str()) {
Some(format!("{method} uri={}", safe_display(uri, 200)))
} else {
Some(method.to_string())
}
}
"notifications/tools/listChanged"
| "notifications/resources/listChanged"
| "notifications/prompts/listChanged" => Some(method.to_string()),
_ => {
if let Some(s) = params
.get("_meta")
.and_then(|m| m.get("theway_summary"))
.and_then(|v| v.as_str())
{
Some(format!("{method} {}", safe_display(s, 200)))
} else {
Some(method.to_string())
}
}
}
}
#[cfg(test)]
tests_bridge_macro::tests_bridge!("triggers/mcp_notification_hook");