677 lines
23 KiB
Rust
677 lines
23 KiB
Rust
use futures::future::BoxFuture;
|
|
use thiserror::Error;
|
|
|
|
use crate::{
|
|
error::Error as McpError,
|
|
model::{
|
|
CancelledNotification, CancelledNotificationParam, JsonRpcMessage, Message, RequestId,
|
|
},
|
|
transport::IntoTransport,
|
|
};
|
|
#[cfg(feature = "client")]
|
|
mod client;
|
|
#[cfg(feature = "client")]
|
|
pub use client::*;
|
|
#[cfg(feature = "server")]
|
|
mod server;
|
|
#[cfg(feature = "server")]
|
|
pub use server::*;
|
|
#[cfg(feature = "tower")]
|
|
mod tower;
|
|
use tokio_util::sync::CancellationToken;
|
|
#[cfg(feature = "tower")]
|
|
pub use tower::*;
|
|
use tracing::instrument;
|
|
#[derive(Error, Debug)]
|
|
#[non_exhaustive]
|
|
pub enum ServiceError {
|
|
#[error("Mcp error: {0}")]
|
|
McpError(McpError),
|
|
#[error("Transport error: {0}")]
|
|
Transport(std::io::Error),
|
|
#[error("Unexpected response type")]
|
|
UnexpectedResponse,
|
|
#[error("task cancelled for reason {}", reason.as_deref().unwrap_or("<unknown>"))]
|
|
Cancelled { reason: Option<String> },
|
|
#[error("request timeout after {}", chrono::Duration::from_std(*timeout).unwrap_or_default())]
|
|
Timeout { timeout: Duration },
|
|
}
|
|
|
|
impl ServiceError {}
|
|
trait TransferObject:
|
|
std::fmt::Debug + Clone + serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static
|
|
{
|
|
}
|
|
|
|
impl<T> TransferObject for T where
|
|
T: std::fmt::Debug
|
|
+ serde::Serialize
|
|
+ serde::de::DeserializeOwned
|
|
+ Send
|
|
+ Sync
|
|
+ 'static
|
|
+ Clone
|
|
{
|
|
}
|
|
|
|
#[allow(private_bounds, reason = "there's no the third implementation")]
|
|
pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
|
|
type Req: TransferObject;
|
|
type Resp: TransferObject;
|
|
type Not: TryInto<CancelledNotification, Error = Self::Not>
|
|
+ From<CancelledNotification>
|
|
+ TransferObject;
|
|
type PeerReq: TransferObject;
|
|
type PeerResp: TransferObject;
|
|
type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
|
|
+ From<CancelledNotification>
|
|
+ TransferObject;
|
|
const IS_CLIENT: bool;
|
|
type Info: TransferObject;
|
|
type PeerInfo: TransferObject;
|
|
}
|
|
|
|
pub type TxJsonRpcMessage<R> =
|
|
JsonRpcMessage<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
|
|
pub type RxJsonRpcMessage<R> = JsonRpcMessage<
|
|
<R as ServiceRole>::PeerReq,
|
|
<R as ServiceRole>::PeerResp,
|
|
<R as ServiceRole>::PeerNot,
|
|
>;
|
|
|
|
pub type TxMessage<R> =
|
|
Message<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
|
|
pub type RxMessage<R> =
|
|
Message<<R as ServiceRole>::PeerReq, <R as ServiceRole>::PeerResp, <R as ServiceRole>::PeerNot>;
|
|
|
|
pub trait Service<R: ServiceRole>: Send + Sync + 'static {
|
|
fn handle_request(
|
|
&self,
|
|
request: R::PeerReq,
|
|
context: RequestContext<R>,
|
|
) -> impl Future<Output = Result<R::Resp, McpError>> + Send + '_;
|
|
fn handle_notification(
|
|
&self,
|
|
notification: R::PeerNot,
|
|
) -> impl Future<Output = Result<(), McpError>> + Send + '_;
|
|
fn get_peer(&self) -> Option<Peer<R>>;
|
|
fn set_peer(&mut self, peer: Peer<R>);
|
|
fn get_info(&self) -> R::Info;
|
|
}
|
|
|
|
pub trait ServiceExt<R: ServiceRole>: Service<R> + Sized {
|
|
/// Convert this service to a dynamic boxed service
|
|
///
|
|
/// This could be very helpful when you want to store the services in a collection
|
|
fn into_dyn(self) -> Box<dyn DynService<R>> {
|
|
Box::new(self)
|
|
}
|
|
fn serve<T, E, A>(
|
|
self,
|
|
transport: T,
|
|
) -> impl Future<Output = Result<RunningService<R, Self>, E>> + Send
|
|
where
|
|
T: IntoTransport<R, E, A>,
|
|
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
|
Self: Sized,
|
|
{
|
|
Self::serve_with_ct(self, transport, Default::default())
|
|
}
|
|
fn serve_with_ct<T, E, A>(
|
|
self,
|
|
transport: T,
|
|
ct: CancellationToken,
|
|
) -> impl Future<Output = Result<RunningService<R, Self>, E>> + Send
|
|
where
|
|
T: IntoTransport<R, E, A>,
|
|
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
|
Self: Sized;
|
|
}
|
|
|
|
impl<R: ServiceRole> Service<R> for Box<dyn DynService<R>> {
|
|
fn handle_request(
|
|
&self,
|
|
request: R::PeerReq,
|
|
context: RequestContext<R>,
|
|
) -> impl Future<Output = Result<R::Resp, McpError>> + Send + '_ {
|
|
DynService::handle_request(self.as_ref(), request, context)
|
|
}
|
|
|
|
fn handle_notification(
|
|
&self,
|
|
notification: R::PeerNot,
|
|
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
|
DynService::handle_notification(self.as_ref(), notification)
|
|
}
|
|
|
|
fn get_peer(&self) -> Option<Peer<R>> {
|
|
DynService::get_peer(self.as_ref())
|
|
}
|
|
|
|
fn set_peer(&mut self, peer: Peer<R>) {
|
|
DynService::set_peer(self.as_mut(), peer)
|
|
}
|
|
|
|
fn get_info(&self) -> R::Info {
|
|
DynService::get_info(self.as_ref())
|
|
}
|
|
}
|
|
|
|
pub trait DynService<R: ServiceRole>: Send + Sync {
|
|
fn handle_request(
|
|
&self,
|
|
request: R::PeerReq,
|
|
context: RequestContext<R>,
|
|
) -> BoxFuture<Result<R::Resp, McpError>>;
|
|
fn handle_notification(&self, notification: R::PeerNot) -> BoxFuture<Result<(), McpError>>;
|
|
fn get_peer(&self) -> Option<Peer<R>>;
|
|
fn set_peer(&mut self, peer: Peer<R>);
|
|
fn get_info(&self) -> R::Info;
|
|
}
|
|
|
|
impl<R: ServiceRole, S: Service<R>> DynService<R> for S {
|
|
fn handle_request(
|
|
&self,
|
|
request: R::PeerReq,
|
|
context: RequestContext<R>,
|
|
) -> BoxFuture<Result<R::Resp, McpError>> {
|
|
Box::pin(self.handle_request(request, context))
|
|
}
|
|
fn handle_notification(&self, notification: R::PeerNot) -> BoxFuture<Result<(), McpError>> {
|
|
Box::pin(self.handle_notification(notification))
|
|
}
|
|
fn get_peer(&self) -> Option<Peer<R>> {
|
|
self.get_peer()
|
|
}
|
|
fn set_peer(&mut self, peer: Peer<R>) {
|
|
self.set_peer(peer)
|
|
}
|
|
fn get_info(&self) -> R::Info {
|
|
self.get_info()
|
|
}
|
|
}
|
|
|
|
use std::{
|
|
collections::HashMap,
|
|
ops::Deref,
|
|
sync::{Arc, atomic::AtomicU32},
|
|
time::Duration,
|
|
};
|
|
|
|
use tokio::sync::mpsc;
|
|
|
|
pub trait RequestIdProvider: Send + Sync + 'static {
|
|
fn next_request_id(&self) -> RequestId;
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct AtomicU32RequestIdProvider {
|
|
id: AtomicU32,
|
|
}
|
|
|
|
impl RequestIdProvider for AtomicU32RequestIdProvider {
|
|
fn next_request_id(&self) -> RequestId {
|
|
RequestId::Number(self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst))
|
|
}
|
|
}
|
|
|
|
type Responder<T> = tokio::sync::oneshot::Sender<T>;
|
|
|
|
/// A handle to a remote request
|
|
///
|
|
/// You can cancel it by call [`RequestHandle::cancel`] with a reason,
|
|
///
|
|
/// or wait for response by call [`RequestHandle::await_response`]
|
|
#[derive(Debug)]
|
|
pub struct RequestHandle<R: ServiceRole> {
|
|
pub rx: tokio::sync::oneshot::Receiver<Result<R::PeerResp, ServiceError>>,
|
|
pub options: PeerRequestOptions,
|
|
pub peer: Peer<R>,
|
|
pub id: RequestId,
|
|
}
|
|
|
|
impl<R: ServiceRole> RequestHandle<R> {
|
|
pub const REQUEST_TIMEOUT_REASON: &str = "request timeout";
|
|
pub async fn await_response(self) -> Result<R::PeerResp, ServiceError> {
|
|
if let Some(timeout) = self.options.timeout {
|
|
let timeout_result = tokio::time::timeout(timeout, async move {
|
|
self.rx
|
|
.await
|
|
.map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
|
|
})
|
|
.await;
|
|
match timeout_result {
|
|
Ok(response) => response,
|
|
Err(_) => {
|
|
let error = Err(ServiceError::Timeout { timeout });
|
|
// cancel this request
|
|
let notification = CancelledNotification {
|
|
params: CancelledNotificationParam {
|
|
request_id: self.id,
|
|
reason: Some(Self::REQUEST_TIMEOUT_REASON.to_owned()),
|
|
},
|
|
method: crate::model::CancelledNotificationMethod,
|
|
};
|
|
let _ = self.peer.send_notification(notification.into()).await;
|
|
error
|
|
}
|
|
}
|
|
} else {
|
|
self.rx
|
|
.await
|
|
.map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
|
|
}
|
|
}
|
|
|
|
/// Cancel this request
|
|
pub async fn cancel(self, reason: Option<String>) -> Result<(), ServiceError> {
|
|
let notification = CancelledNotification {
|
|
params: CancelledNotificationParam {
|
|
request_id: self.id,
|
|
reason,
|
|
},
|
|
method: crate::model::CancelledNotificationMethod,
|
|
};
|
|
self.peer.send_notification(notification.into()).await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum PeerSinkMessage<R: ServiceRole> {
|
|
Request(
|
|
R::Req,
|
|
RequestId,
|
|
Responder<Result<R::PeerResp, ServiceError>>,
|
|
),
|
|
Notification(R::Not, Responder<Result<(), ServiceError>>),
|
|
}
|
|
|
|
/// An interface to fetch the remote client or server
|
|
///
|
|
/// For general purpose, call [`Peer::send_request`] or [`Peer::send_notification`] to send message to remote peer.
|
|
///
|
|
/// To create a cancellable request, call [`Peer::send_cancellable_request`].
|
|
#[derive(Clone)]
|
|
pub struct Peer<R: ServiceRole> {
|
|
tx: mpsc::Sender<PeerSinkMessage<R>>,
|
|
request_id_provider: Arc<dyn RequestIdProvider>,
|
|
info: Arc<R::PeerInfo>,
|
|
}
|
|
|
|
impl<R: ServiceRole> std::fmt::Debug for Peer<R> {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("PeerSink")
|
|
.field("tx", &self.tx)
|
|
.field("is_client", &R::IS_CLIENT)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
type ProxyOutbound<R> = mpsc::Receiver<PeerSinkMessage<R>>;
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct PeerRequestOptions {
|
|
timeout: Option<Duration>,
|
|
}
|
|
|
|
impl PeerRequestOptions {
|
|
pub fn no_options() -> Self {
|
|
Self::default()
|
|
}
|
|
}
|
|
|
|
impl<R: ServiceRole> Peer<R> {
|
|
const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
|
|
pub fn new(
|
|
request_id_provider: Arc<dyn RequestIdProvider>,
|
|
peer_info: R::PeerInfo,
|
|
) -> (Peer<R>, ProxyOutbound<R>) {
|
|
let (tx, rx) = mpsc::channel(Self::CLIENT_CHANNEL_BUFFER_SIZE);
|
|
(
|
|
Self {
|
|
tx,
|
|
request_id_provider,
|
|
info: peer_info.into(),
|
|
},
|
|
rx,
|
|
)
|
|
}
|
|
pub async fn send_notification(&self, notification: R::Not) -> Result<(), ServiceError> {
|
|
let (responder, receiver) = tokio::sync::oneshot::channel();
|
|
self.tx
|
|
.send(PeerSinkMessage::Notification(notification, responder))
|
|
.await
|
|
.map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
|
|
receiver
|
|
.await
|
|
.map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
|
|
}
|
|
pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
|
|
self.send_cancellable_request(request, PeerRequestOptions::no_options())
|
|
.await?
|
|
.await_response()
|
|
.await
|
|
}
|
|
pub async fn send_cancellable_request(
|
|
&self,
|
|
request: R::Req,
|
|
options: PeerRequestOptions,
|
|
) -> Result<RequestHandle<R>, ServiceError> {
|
|
let id = self.request_id_provider.next_request_id();
|
|
let (responder, receiver) = tokio::sync::oneshot::channel();
|
|
self.tx
|
|
.send(PeerSinkMessage::Request(request, id.clone(), responder))
|
|
.await
|
|
.map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
|
|
Ok(RequestHandle {
|
|
id,
|
|
rx: receiver,
|
|
options,
|
|
peer: self.clone(),
|
|
})
|
|
}
|
|
pub fn peer_info(&self) -> &R::PeerInfo {
|
|
&self.info
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct RunningService<R: ServiceRole, S: Service<R>> {
|
|
service: Arc<S>,
|
|
peer: Peer<R>,
|
|
handle: tokio::task::JoinHandle<QuitReason>,
|
|
/// cancellation token
|
|
ct: CancellationToken,
|
|
}
|
|
impl<R: ServiceRole, S: Service<R>> Deref for RunningService<R, S> {
|
|
type Target = Peer<R>;
|
|
|
|
fn deref(&self) -> &Self::Target {
|
|
self.peer()
|
|
}
|
|
}
|
|
|
|
impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
|
|
#[inline]
|
|
pub fn peer(&self) -> &Peer<R> {
|
|
&self.peer
|
|
}
|
|
#[inline]
|
|
pub fn service(&self) -> &S {
|
|
self.service.as_ref()
|
|
}
|
|
pub async fn waiting(self) -> Result<QuitReason, tokio::task::JoinError> {
|
|
self.handle.await
|
|
}
|
|
pub async fn cancel(self) -> Result<QuitReason, tokio::task::JoinError> {
|
|
self.ct.cancel();
|
|
self.waiting().await
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
|
|
pub enum QuitReason {
|
|
Cancelled,
|
|
Closed,
|
|
}
|
|
|
|
/// Request execution context
|
|
#[derive(Debug, Clone)]
|
|
pub struct RequestContext<R: ServiceRole> {
|
|
/// this token will be cancelled when the [`CancelledNotification`] is received.
|
|
pub ct: CancellationToken,
|
|
pub id: RequestId,
|
|
/// An interface to fetch the remote client or server
|
|
pub peer: Peer<R>,
|
|
}
|
|
|
|
/// Use this function to skip initialization process
|
|
pub async fn serve_directly<R, S, T, E, A>(
|
|
service: S,
|
|
transport: T,
|
|
peer_info: R::PeerInfo,
|
|
) -> Result<RunningService<R, S>, E>
|
|
where
|
|
R: ServiceRole,
|
|
S: Service<R>,
|
|
T: IntoTransport<R, E, A>,
|
|
E: std::error::Error + Send + Sync + 'static,
|
|
{
|
|
serve_directly_with_ct(service, transport, peer_info, Default::default()).await
|
|
}
|
|
|
|
/// Use this function to skip initialization process
|
|
pub async fn serve_directly_with_ct<R, S, T, E, A>(
|
|
service: S,
|
|
transport: T,
|
|
peer_info: R::PeerInfo,
|
|
ct: CancellationToken,
|
|
) -> Result<RunningService<R, S>, E>
|
|
where
|
|
R: ServiceRole,
|
|
S: Service<R>,
|
|
T: IntoTransport<R, E, A>,
|
|
E: std::error::Error + Send + Sync + 'static,
|
|
{
|
|
serve_inner(service, transport, peer_info, Default::default(), ct).await
|
|
}
|
|
|
|
#[instrument(skip_all)]
|
|
async fn serve_inner<R, S, T, E, A>(
|
|
mut service: S,
|
|
transport: T,
|
|
peer_info: R::PeerInfo,
|
|
id_provider: Arc<AtomicU32RequestIdProvider>,
|
|
ct: CancellationToken,
|
|
) -> Result<RunningService<R, S>, E>
|
|
where
|
|
R: ServiceRole,
|
|
S: Service<R>,
|
|
T: IntoTransport<R, E, A>,
|
|
E: std::error::Error + Send + Sync + 'static,
|
|
{
|
|
use futures::{SinkExt, StreamExt};
|
|
const SINK_PROXY_BUFFER_SIZE: usize = 64;
|
|
let (sink_proxy_tx, mut sink_proxy_rx) = tokio::sync::mpsc::channel::<
|
|
Message<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>,
|
|
>(SINK_PROXY_BUFFER_SIZE);
|
|
|
|
if R::IS_CLIENT {
|
|
tracing::info!(?peer_info, "Service initialized as client");
|
|
} else {
|
|
tracing::info!(?peer_info, "Service initialized as server");
|
|
}
|
|
|
|
let (peer, mut peer_proxy) = <Peer<R>>::new(id_provider, peer_info);
|
|
service.set_peer(peer.clone());
|
|
let mut local_responder_pool = HashMap::new();
|
|
let mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
|
|
let shared_service = Arc::new(service);
|
|
// for return
|
|
let service = shared_service.clone();
|
|
|
|
// let message_sink = tokio::sync::
|
|
// let mut stream = std::pin::pin!(stream);
|
|
let serve_loop_ct = ct.child_token();
|
|
let peer_return: Peer<R> = peer.clone();
|
|
let handle = tokio::spawn(async move {
|
|
let (mut sink, mut stream) = transport.into_transport();
|
|
let mut sink = std::pin::pin!(sink);
|
|
let mut stream = std::pin::pin!(stream);
|
|
#[derive(Debug)]
|
|
enum Event<P, R, T> {
|
|
ProxyMessage(P),
|
|
PeerMessage(R),
|
|
ToSink(T),
|
|
}
|
|
let quit_reason = loop {
|
|
let evt = tokio::select! {
|
|
m = sink_proxy_rx.recv() => {
|
|
if let Some(m) = m {
|
|
Event::ToSink(m)
|
|
} else {
|
|
continue
|
|
}
|
|
}
|
|
m = stream.next() => {
|
|
if let Some(m) = m {
|
|
Event::PeerMessage(m.into_message())
|
|
} else {
|
|
// input stream closed
|
|
tracing::info!("input stream terminated");
|
|
break QuitReason::Closed
|
|
}
|
|
}
|
|
m = peer_proxy.recv() => {
|
|
if let Some(m) = m {
|
|
Event::ProxyMessage(m)
|
|
} else {
|
|
continue
|
|
}
|
|
}
|
|
_ = serve_loop_ct.cancelled() => {
|
|
tracing::info!("task cancelled");
|
|
break QuitReason::Cancelled
|
|
}
|
|
};
|
|
tracing::debug!(?evt, "new event");
|
|
match evt {
|
|
// response and error
|
|
Event::ToSink(e) => {
|
|
if let Some(id) = match &e {
|
|
Message::Response(_, id) => Some(id),
|
|
Message::Error(_, id) => Some(id),
|
|
_ => None,
|
|
} {
|
|
if let Some(ct) = local_ct_pool.remove(id) {
|
|
ct.cancel();
|
|
}
|
|
let send_result = sink.send(e.into_json_rpc_message()).await;
|
|
if let Err(error) = send_result {
|
|
tracing::error!(%error, "fail to response message");
|
|
}
|
|
}
|
|
}
|
|
Event::ProxyMessage(PeerSinkMessage::Request(request, id, responder)) => {
|
|
local_responder_pool.insert(id.clone(), responder);
|
|
let send_result = sink
|
|
.send(Message::Request(request, id.clone()).into_json_rpc_message())
|
|
.await;
|
|
if let Err(e) = send_result {
|
|
if let Some(responder) = local_responder_pool.remove(&id) {
|
|
let _ = responder
|
|
.send(Err(ServiceError::Transport(std::io::Error::other(e))));
|
|
}
|
|
}
|
|
}
|
|
Event::ProxyMessage(PeerSinkMessage::Notification(notification, responder)) => {
|
|
// catch cancellation notification
|
|
let mut cancellation_param = None;
|
|
let notification = match notification.try_into() {
|
|
Ok::<CancelledNotification, _>(cancelled) => {
|
|
cancellation_param.replace(cancelled.params.clone());
|
|
cancelled.into()
|
|
}
|
|
Err(notification) => notification,
|
|
};
|
|
let send_result = sink
|
|
.send(Message::Notification(notification).into_json_rpc_message())
|
|
.await;
|
|
if let Err(e) = send_result {
|
|
let _ =
|
|
responder.send(Err(ServiceError::Transport(std::io::Error::other(e))));
|
|
}
|
|
if let Some(param) = cancellation_param {
|
|
if let Some(responder) = local_responder_pool.remove(¶m.request_id) {
|
|
tracing::info!(id = %param.request_id, reason = param.reason, "cancelled");
|
|
let _response_result = responder.send(Err(ServiceError::Cancelled {
|
|
reason: param.reason.clone(),
|
|
}));
|
|
}
|
|
}
|
|
}
|
|
Event::PeerMessage(Message::Request(request, id)) => {
|
|
tracing::info!(%id, ?request, "received request");
|
|
{
|
|
let service = shared_service.clone();
|
|
let sink = sink_proxy_tx.clone();
|
|
let request_ct = serve_loop_ct.child_token();
|
|
let context_ct = request_ct.child_token();
|
|
local_ct_pool.insert(id.clone(), request_ct);
|
|
let context = RequestContext {
|
|
ct: context_ct,
|
|
id: id.clone(),
|
|
peer: peer.clone(),
|
|
};
|
|
tokio::spawn(async move {
|
|
let result = service.handle_request(request, context).await;
|
|
let response = match result {
|
|
Ok(result) => {
|
|
tracing::info!(%id, ?result, "response message");
|
|
Message::Response(result, id)
|
|
}
|
|
Err(error) => {
|
|
tracing::warn!(%id, ?error, "response error");
|
|
Message::Error(error, id)
|
|
}
|
|
};
|
|
let _send_result = sink.send(response).await;
|
|
});
|
|
}
|
|
}
|
|
Event::PeerMessage(Message::Notification(notification)) => {
|
|
tracing::info!(?notification, "received notification");
|
|
// catch cancelled notification
|
|
let notification = match notification.try_into() {
|
|
Ok::<CancelledNotification, _>(cancelled) => {
|
|
if let Some(ct) = local_ct_pool.remove(&cancelled.params.request_id) {
|
|
tracing::info!(id = %cancelled.params.request_id, reason = cancelled.params.reason, "cancelled");
|
|
ct.cancel();
|
|
}
|
|
cancelled.into()
|
|
}
|
|
Err(notification) => notification,
|
|
};
|
|
{
|
|
let service = shared_service.clone();
|
|
tokio::spawn(async move {
|
|
let result = service.handle_notification(notification).await;
|
|
if let Err(error) = result {
|
|
tracing::warn!(%error, "Error sending notification");
|
|
}
|
|
});
|
|
}
|
|
}
|
|
Event::PeerMessage(Message::Response(result, id)) => {
|
|
if let Some(responder) = local_responder_pool.remove(&id) {
|
|
let response_result = responder.send(Ok(result));
|
|
if let Err(_error) = response_result {
|
|
tracing::warn!(%id, "Error sending response");
|
|
}
|
|
}
|
|
}
|
|
Event::PeerMessage(Message::Error(error, id)) => {
|
|
if let Some(responder) = local_responder_pool.remove(&id) {
|
|
let _response_result = responder.send(Err(ServiceError::McpError(error)));
|
|
if let Err(_error) = _response_result {
|
|
tracing::warn!(%id, "Error sending response");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
let sink_close_result = sink.close().await;
|
|
if let Err(e) = sink_close_result {
|
|
tracing::error!(%e, "fail to close sink");
|
|
}
|
|
tracing::info!(?quit_reason, "serve finished");
|
|
quit_reason
|
|
});
|
|
Ok(RunningService {
|
|
service,
|
|
peer: peer_return,
|
|
handle,
|
|
ct,
|
|
})
|
|
}
|