use async_trait::async_trait;
use std::sync::Arc;
use tracing::{debug, error, info};
use crate::StreamManager;
use turul_mcp_json_rpc_server::JsonRpcNotification;
use turul_mcp_protocol::notifications::{
CancelledNotification, LoggingMessageNotification, ProgressNotification,
PromptListChangedNotification, ResourceListChangedNotification, ResourceUpdatedNotification,
ToolListChangedNotification,
};
#[async_trait]
pub trait NotificationBroadcaster: Send + Sync {
async fn send_progress_notification(
&self,
session_id: &str,
notification: ProgressNotification,
) -> Result<(), BroadcastError>;
async fn send_message_notification(
&self,
session_id: &str,
notification: LoggingMessageNotification,
) -> Result<(), BroadcastError>;
async fn send_resource_updated_notification(
&self,
session_id: &str,
notification: ResourceUpdatedNotification,
) -> Result<(), BroadcastError>;
async fn send_resource_list_changed_notification(
&self,
session_id: &str,
notification: ResourceListChangedNotification,
) -> Result<(), BroadcastError>;
async fn send_tool_list_changed_notification(
&self,
session_id: &str,
notification: ToolListChangedNotification,
) -> Result<(), BroadcastError>;
async fn send_prompt_list_changed_notification(
&self,
session_id: &str,
notification: PromptListChangedNotification,
) -> Result<(), BroadcastError>;
async fn send_cancelled_notification(
&self,
session_id: &str,
notification: CancelledNotification,
) -> Result<(), BroadcastError>;
async fn broadcast_to_all_sessions(
&self,
notification: JsonRpcNotification,
) -> Result<Vec<String>, BroadcastError>;
async fn send_notification(
&self,
session_id: &str,
notification: JsonRpcNotification,
) -> Result<(), BroadcastError>;
}
#[derive(Debug, thiserror::Error)]
pub enum BroadcastError {
#[error("Session not found: {0}")]
SessionNotFound(String),
#[error("Broadcasting failed: {0}")]
BroadcastFailed(String),
#[error("Serialization error: {0}")]
SerializationError(#[from] serde_json::Error),
}
pub struct StreamManagerNotificationBroadcaster {
stream_manager: Arc<StreamManager>,
}
impl StreamManagerNotificationBroadcaster {
pub fn new(stream_manager: Arc<StreamManager>) -> Self {
Self { stream_manager }
}
}
pub mod conversion {
use super::*;
use std::collections::HashMap;
pub fn progress_to_json_rpc(notification: ProgressNotification) -> JsonRpcNotification {
let mut params = HashMap::new();
params.insert(
"progressToken".to_string(),
serde_json::json!(notification.params.progress_token),
);
params.insert(
"progress".to_string(),
serde_json::json!(notification.params.progress),
);
if let Some(total) = notification.params.total {
params.insert("total".to_string(), serde_json::json!(total));
}
if let Some(message) = notification.params.message {
params.insert("message".to_string(), serde_json::json!(message));
}
if let Some(meta) = notification.params.meta {
params.insert("_meta".to_string(), serde_json::json!(meta));
}
JsonRpcNotification::new_with_object_params(notification.method, params)
}
pub fn message_to_json_rpc(notification: LoggingMessageNotification) -> JsonRpcNotification {
let mut params = HashMap::new();
params.insert(
"level".to_string(),
serde_json::json!(notification.params.level),
);
params.insert("data".to_string(), notification.params.data);
if let Some(logger) = notification.params.logger {
params.insert("logger".to_string(), serde_json::json!(logger));
}
if let Some(meta) = notification.params.meta {
params.insert("_meta".to_string(), serde_json::json!(meta));
}
JsonRpcNotification::new_with_object_params(notification.method, params)
}
pub fn resource_updated_to_json_rpc(
notification: ResourceUpdatedNotification,
) -> JsonRpcNotification {
let mut params = HashMap::new();
params.insert(
"uri".to_string(),
serde_json::json!(notification.params.uri),
);
if let Some(meta) = notification.params.meta {
params.insert("_meta".to_string(), serde_json::json!(meta));
}
JsonRpcNotification::new_with_object_params(notification.method, params)
}
pub fn resource_list_changed_to_json_rpc(
notification: ResourceListChangedNotification,
) -> JsonRpcNotification {
if let Some(params) = notification.params {
if let Some(meta) = params.meta {
let mut param_map = HashMap::new();
param_map.insert("_meta".to_string(), serde_json::json!(meta));
JsonRpcNotification::new_with_object_params(notification.method, param_map)
} else {
JsonRpcNotification::new_no_params(notification.method)
}
} else {
JsonRpcNotification::new_no_params(notification.method)
}
}
pub fn tool_list_changed_to_json_rpc(
notification: ToolListChangedNotification,
) -> JsonRpcNotification {
if let Some(params) = notification.params {
if let Some(meta) = params.meta {
let mut param_map = HashMap::new();
param_map.insert("_meta".to_string(), serde_json::json!(meta));
JsonRpcNotification::new_with_object_params(notification.method, param_map)
} else {
JsonRpcNotification::new_no_params(notification.method)
}
} else {
JsonRpcNotification::new_no_params(notification.method)
}
}
pub fn prompt_list_changed_to_json_rpc(
notification: PromptListChangedNotification,
) -> JsonRpcNotification {
if let Some(params) = notification.params {
if let Some(meta) = params.meta {
let mut param_map = HashMap::new();
param_map.insert("_meta".to_string(), serde_json::json!(meta));
JsonRpcNotification::new_with_object_params(notification.method, param_map)
} else {
JsonRpcNotification::new_no_params(notification.method)
}
} else {
JsonRpcNotification::new_no_params(notification.method)
}
}
pub fn cancelled_to_json_rpc(notification: CancelledNotification) -> JsonRpcNotification {
let mut params = HashMap::new();
params.insert(
"requestId".to_string(),
serde_json::json!(notification.params.request_id),
);
if let Some(reason) = notification.params.reason {
params.insert("reason".to_string(), serde_json::json!(reason));
}
if let Some(meta) = notification.params.meta {
params.insert("_meta".to_string(), serde_json::json!(meta));
}
JsonRpcNotification::new_with_object_params(notification.method, params)
}
}
#[async_trait]
impl NotificationBroadcaster for StreamManagerNotificationBroadcaster {
async fn send_progress_notification(
&self,
session_id: &str,
notification: ProgressNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::progress_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn send_message_notification(
&self,
session_id: &str,
notification: LoggingMessageNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::message_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn send_resource_updated_notification(
&self,
session_id: &str,
notification: ResourceUpdatedNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::resource_updated_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn send_resource_list_changed_notification(
&self,
session_id: &str,
notification: ResourceListChangedNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::resource_list_changed_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn send_tool_list_changed_notification(
&self,
session_id: &str,
notification: ToolListChangedNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::tool_list_changed_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn send_prompt_list_changed_notification(
&self,
session_id: &str,
notification: PromptListChangedNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::prompt_list_changed_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn send_cancelled_notification(
&self,
session_id: &str,
notification: CancelledNotification,
) -> Result<(), BroadcastError> {
let json_rpc_notification = conversion::cancelled_to_json_rpc(notification);
self.send_notification(session_id, json_rpc_notification)
.await
}
async fn broadcast_to_all_sessions(
&self,
notification: JsonRpcNotification,
) -> Result<Vec<String>, BroadcastError> {
let sse_data =
serde_json::to_value(¬ification).map_err(BroadcastError::SerializationError)?;
match self
.stream_manager
.broadcast_to_all_sessions(
notification.method.clone(), sse_data,
)
.await
{
Ok(failed_sessions) => {
info!(
"📡 Broadcast JSON-RPC notification to all sessions: method={}, failed={}",
notification.method,
failed_sessions.len()
);
Ok(failed_sessions)
}
Err(e) => {
error!(
"❌ Failed to broadcast JSON-RPC notification: method={}, error={}",
notification.method, e
);
Err(BroadcastError::BroadcastFailed(e.to_string()))
}
}
}
async fn send_notification(
&self,
session_id: &str,
notification: JsonRpcNotification,
) -> Result<(), BroadcastError> {
let sse_data =
serde_json::to_value(¬ification).map_err(BroadcastError::SerializationError)?;
match self
.stream_manager
.broadcast_to_session(
session_id,
notification.method.clone(), sse_data,
)
.await
{
Ok(event_id) => {
debug!(
"✅ Sent JSON-RPC notification: session={}, method={}, event_id={}",
session_id, notification.method, event_id
);
Ok(())
}
Err(e) => {
error!(
"❌ Failed to send JSON-RPC notification: session={}, method={}, error={}",
session_id, notification.method, e
);
Err(BroadcastError::BroadcastFailed(e.to_string()))
}
}
}
}
pub type SharedNotificationBroadcaster = Arc<dyn NotificationBroadcaster + Send + Sync>;