95 lines
2.7 KiB
Rust
95 lines
2.7 KiB
Rust
use std::sync::Arc;
|
|
|
|
use rmcp::{
|
|
ClientHandler, ServerHandler, ServiceExt,
|
|
model::{
|
|
ResourceUpdatedNotificationParam, ServerCapabilities, ServerInfo, SubscribeRequestParam,
|
|
},
|
|
};
|
|
use tokio::sync::Notify;
|
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
|
|
|
pub struct Server {}
|
|
|
|
impl ServerHandler for Server {
|
|
fn get_info(&self) -> ServerInfo {
|
|
ServerInfo {
|
|
capabilities: ServerCapabilities::builder()
|
|
.enable_resources()
|
|
.enable_resources_subscribe()
|
|
.enable_resources_list_changed()
|
|
.build(),
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
async fn subscribe(
|
|
&self,
|
|
request: rmcp::model::SubscribeRequestParam,
|
|
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
|
) -> Result<(), rmcp::ErrorData> {
|
|
let uri = request.uri;
|
|
let peer = context.peer;
|
|
|
|
tokio::spawn(async move {
|
|
let span = tracing::info_span!("subscribe", uri = %uri);
|
|
let _enter = span.enter();
|
|
|
|
if let Err(e) = peer
|
|
.notify_resource_updated(ResourceUpdatedNotificationParam { uri: uri.clone() })
|
|
.await
|
|
{
|
|
panic!("Failed to send notification: {}", e);
|
|
}
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
pub struct Client {
|
|
receive_signal: Arc<Notify>,
|
|
}
|
|
|
|
impl ClientHandler for Client {
|
|
async fn on_resource_updated(
|
|
&self,
|
|
params: rmcp::model::ResourceUpdatedNotificationParam,
|
|
_context: rmcp::service::NotificationContext<rmcp::RoleClient>,
|
|
) {
|
|
let uri = params.uri;
|
|
tracing::info!("Resource updated: {}", uri);
|
|
self.receive_signal.notify_one();
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_server_notification() -> anyhow::Result<()> {
|
|
let _ = tracing_subscriber::registry()
|
|
.with(
|
|
tracing_subscriber::EnvFilter::try_from_default_env()
|
|
.unwrap_or_else(|_| "debug".to_string().into()),
|
|
)
|
|
.with(tracing_subscriber::fmt::layer())
|
|
.try_init();
|
|
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
|
tokio::spawn(async move {
|
|
let server = Server {}.serve(server_transport).await?;
|
|
server.waiting().await?;
|
|
anyhow::Ok(())
|
|
});
|
|
let receive_signal = Arc::new(Notify::new());
|
|
let client = Client {
|
|
receive_signal: receive_signal.clone(),
|
|
}
|
|
.serve(client_transport)
|
|
.await?;
|
|
client
|
|
.subscribe(SubscribeRequestParam {
|
|
uri: "test://test-resource".to_owned(),
|
|
})
|
|
.await?;
|
|
receive_signal.notified().await;
|
|
client.cancel().await?;
|
|
Ok(())
|
|
}
|