1049 lines
37 KiB
Rust
1049 lines
37 KiB
Rust
use futures::FutureExt;
|
||
#[cfg(not(feature = "local"))]
|
||
use futures::future::BoxFuture;
|
||
#[cfg(feature = "local")]
|
||
use futures::future::LocalBoxFuture;
|
||
use thiserror::Error;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Conditional Send helpers
|
||
//
|
||
// `MaybeSend` – supertrait alias: `Send + Sync` without `local`, empty with `local`
|
||
// `MaybeSendFuture` – future bound alias: `Send` without `local`, empty with `local`
|
||
// `MaybeBoxFuture` – boxed future type: `BoxFuture` without `local`, `LocalBoxFuture` with `local`
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[cfg(not(feature = "local"))]
|
||
#[doc(hidden)]
|
||
pub trait MaybeSend: Send + Sync {}
|
||
#[cfg(not(feature = "local"))]
|
||
impl<T: Send + Sync> MaybeSend for T {}
|
||
|
||
#[cfg(feature = "local")]
|
||
#[doc(hidden)]
|
||
pub trait MaybeSend {}
|
||
#[cfg(feature = "local")]
|
||
impl<T> MaybeSend for T {}
|
||
|
||
#[cfg(not(feature = "local"))]
|
||
#[doc(hidden)]
|
||
pub trait MaybeSendFuture: Send {}
|
||
#[cfg(not(feature = "local"))]
|
||
impl<T: Send> MaybeSendFuture for T {}
|
||
|
||
#[cfg(feature = "local")]
|
||
#[doc(hidden)]
|
||
pub trait MaybeSendFuture {}
|
||
#[cfg(feature = "local")]
|
||
impl<T> MaybeSendFuture for T {}
|
||
|
||
#[cfg(not(feature = "local"))]
|
||
pub(crate) type MaybeBoxFuture<'a, T> = BoxFuture<'a, T>;
|
||
#[cfg(feature = "local")]
|
||
pub(crate) type MaybeBoxFuture<'a, T> = LocalBoxFuture<'a, T>;
|
||
|
||
#[cfg(feature = "server")]
|
||
use crate::model::ServerJsonRpcMessage;
|
||
use crate::{
|
||
error::ErrorData as McpError,
|
||
model::{
|
||
CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta,
|
||
JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Meta,
|
||
NumberOrString, ProgressToken, RequestId,
|
||
},
|
||
transport::{DynamicTransportError, IntoTransport, Transport},
|
||
};
|
||
#[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, DropGuard};
|
||
#[cfg(feature = "tower")]
|
||
pub use tower::*;
|
||
use tracing::{Instrument as _, instrument};
|
||
#[derive(Error, Debug)]
|
||
#[non_exhaustive]
|
||
pub enum ServiceError {
|
||
#[error("Mcp error: {0}")]
|
||
McpError(McpError),
|
||
#[error("Transport send error: {0}")]
|
||
TransportSend(DynamicTransportError),
|
||
#[error("Transport closed")]
|
||
TransportClosed,
|
||
#[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 },
|
||
}
|
||
|
||
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 + GetMeta + GetExtensions;
|
||
type Resp: TransferObject;
|
||
type Not: TryInto<CancelledNotification, Error = Self::Not>
|
||
+ From<CancelledNotification>
|
||
+ TransferObject;
|
||
type PeerReq: TransferObject + GetMeta + GetExtensions;
|
||
type PeerResp: TransferObject;
|
||
type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
|
||
+ From<CancelledNotification>
|
||
+ TransferObject
|
||
+ GetMeta
|
||
+ GetExtensions;
|
||
type InitializeError;
|
||
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,
|
||
>;
|
||
|
||
#[cfg(not(feature = "local"))]
|
||
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>> + MaybeSendFuture + '_;
|
||
fn handle_notification(
|
||
&self,
|
||
notification: R::PeerNot,
|
||
context: NotificationContext<R>,
|
||
) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_;
|
||
fn get_info(&self) -> R::Info;
|
||
}
|
||
|
||
#[cfg(feature = "local")]
|
||
pub trait Service<R: ServiceRole>: 'static {
|
||
fn handle_request(
|
||
&self,
|
||
request: R::PeerReq,
|
||
context: RequestContext<R>,
|
||
) -> impl Future<Output = Result<R::Resp, McpError>> + MaybeSendFuture + '_;
|
||
fn handle_notification(
|
||
&self,
|
||
notification: R::PeerNot,
|
||
context: NotificationContext<R>,
|
||
) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_;
|
||
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>, R::InitializeError>> + MaybeSendFuture
|
||
where
|
||
T: IntoTransport<R, E, A>,
|
||
E: std::error::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>, R::InitializeError>> + MaybeSendFuture
|
||
where
|
||
T: IntoTransport<R, E, A>,
|
||
E: std::error::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>> + MaybeSendFuture + '_ {
|
||
DynService::handle_request(self.as_ref(), request, context)
|
||
}
|
||
|
||
fn handle_notification(
|
||
&self,
|
||
notification: R::PeerNot,
|
||
context: NotificationContext<R>,
|
||
) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_ {
|
||
DynService::handle_notification(self.as_ref(), notification, context)
|
||
}
|
||
|
||
fn get_info(&self) -> R::Info {
|
||
DynService::get_info(self.as_ref())
|
||
}
|
||
}
|
||
|
||
#[cfg(not(feature = "local"))]
|
||
pub trait DynService<R: ServiceRole>: Send + Sync {
|
||
fn handle_request(
|
||
&self,
|
||
request: R::PeerReq,
|
||
context: RequestContext<R>,
|
||
) -> MaybeBoxFuture<'_, Result<R::Resp, McpError>>;
|
||
fn handle_notification(
|
||
&self,
|
||
notification: R::PeerNot,
|
||
context: NotificationContext<R>,
|
||
) -> MaybeBoxFuture<'_, Result<(), McpError>>;
|
||
fn get_info(&self) -> R::Info;
|
||
}
|
||
|
||
#[cfg(feature = "local")]
|
||
pub trait DynService<R: ServiceRole> {
|
||
fn handle_request(
|
||
&self,
|
||
request: R::PeerReq,
|
||
context: RequestContext<R>,
|
||
) -> MaybeBoxFuture<'_, Result<R::Resp, McpError>>;
|
||
fn handle_notification(
|
||
&self,
|
||
notification: R::PeerNot,
|
||
context: NotificationContext<R>,
|
||
) -> MaybeBoxFuture<'_, Result<(), McpError>>;
|
||
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>,
|
||
) -> MaybeBoxFuture<'_, Result<R::Resp, McpError>> {
|
||
Box::pin(self.handle_request(request, context))
|
||
}
|
||
fn handle_notification(
|
||
&self,
|
||
notification: R::PeerNot,
|
||
context: NotificationContext<R>,
|
||
) -> MaybeBoxFuture<'_, Result<(), McpError>> {
|
||
Box::pin(self.handle_notification(notification, context))
|
||
}
|
||
fn get_info(&self) -> R::Info {
|
||
self.get_info()
|
||
}
|
||
}
|
||
|
||
use std::{
|
||
collections::{HashMap, VecDeque},
|
||
ops::Deref,
|
||
sync::{Arc, atomic::AtomicU64},
|
||
time::Duration,
|
||
};
|
||
|
||
use tokio::sync::mpsc;
|
||
|
||
pub trait RequestIdProvider: Send + Sync + 'static {
|
||
fn next_request_id(&self) -> RequestId;
|
||
}
|
||
|
||
pub trait ProgressTokenProvider: Send + Sync + 'static {
|
||
fn next_progress_token(&self) -> ProgressToken;
|
||
}
|
||
|
||
pub type AtomicU32RequestIdProvider = AtomicU32Provider;
|
||
pub type AtomicU32ProgressTokenProvider = AtomicU32Provider;
|
||
|
||
#[derive(Debug, Default)]
|
||
pub struct AtomicU32Provider {
|
||
id: AtomicU64,
|
||
}
|
||
|
||
impl RequestIdProvider for AtomicU32Provider {
|
||
fn next_request_id(&self) -> RequestId {
|
||
let id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||
// Safe conversion: we start at 0 and increment by 1, so we won't overflow i64::MAX in practice
|
||
RequestId::Number(id as i64)
|
||
}
|
||
}
|
||
|
||
impl ProgressTokenProvider for AtomicU32Provider {
|
||
fn next_progress_token(&self) -> ProgressToken {
|
||
let id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||
ProgressToken(NumberOrString::Number(id as i64))
|
||
}
|
||
}
|
||
|
||
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,
|
||
pub progress_token: ProgressToken,
|
||
}
|
||
|
||
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::TransportClosed)?
|
||
})
|
||
.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,
|
||
extensions: Default::default(),
|
||
};
|
||
let _ = self.peer.send_notification(notification.into()).await;
|
||
error
|
||
}
|
||
}
|
||
} else {
|
||
self.rx.await.map_err(|_e| ServiceError::TransportClosed)?
|
||
}
|
||
}
|
||
|
||
/// 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,
|
||
extensions: Default::default(),
|
||
};
|
||
self.peer.send_notification(notification.into()).await?;
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub(crate) enum PeerSinkMessage<R: ServiceRole> {
|
||
Request {
|
||
request: R::Req,
|
||
id: RequestId,
|
||
responder: Responder<Result<R::PeerResp, ServiceError>>,
|
||
},
|
||
Notification {
|
||
notification: R::Not,
|
||
responder: 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_request_with_option`].
|
||
#[derive(Clone)]
|
||
pub struct Peer<R: ServiceRole> {
|
||
tx: mpsc::Sender<PeerSinkMessage<R>>,
|
||
request_id_provider: Arc<dyn RequestIdProvider>,
|
||
progress_token_provider: Arc<dyn ProgressTokenProvider>,
|
||
info: Arc<tokio::sync::OnceCell<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 {
|
||
pub timeout: Option<Duration>,
|
||
pub meta: Option<Meta>,
|
||
}
|
||
|
||
impl PeerRequestOptions {
|
||
pub fn no_options() -> Self {
|
||
Self::default()
|
||
}
|
||
}
|
||
|
||
impl<R: ServiceRole> Peer<R> {
|
||
const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
|
||
pub(crate) fn new(
|
||
request_id_provider: Arc<dyn RequestIdProvider>,
|
||
peer_info: Option<R::PeerInfo>,
|
||
) -> (Peer<R>, ProxyOutbound<R>) {
|
||
let (tx, rx) = mpsc::channel(Self::CLIENT_CHANNEL_BUFFER_SIZE);
|
||
(
|
||
Self {
|
||
tx,
|
||
request_id_provider,
|
||
progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
|
||
info: Arc::new(tokio::sync::OnceCell::new_with(peer_info)),
|
||
},
|
||
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::TransportClosed)?;
|
||
receiver.await.map_err(|_e| ServiceError::TransportClosed)?
|
||
}
|
||
pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
|
||
self.send_request_with_option(request, PeerRequestOptions::no_options())
|
||
.await?
|
||
.await_response()
|
||
.await
|
||
}
|
||
|
||
pub async fn send_cancellable_request(
|
||
&self,
|
||
request: R::Req,
|
||
options: PeerRequestOptions,
|
||
) -> Result<RequestHandle<R>, ServiceError> {
|
||
self.send_request_with_option(request, options).await
|
||
}
|
||
|
||
pub async fn send_request_with_option(
|
||
&self,
|
||
mut request: R::Req,
|
||
options: PeerRequestOptions,
|
||
) -> Result<RequestHandle<R>, ServiceError> {
|
||
let id = self.request_id_provider.next_request_id();
|
||
let progress_token = self.progress_token_provider.next_progress_token();
|
||
request
|
||
.get_meta_mut()
|
||
.set_progress_token(progress_token.clone());
|
||
if let Some(meta) = options.meta.clone() {
|
||
request.get_meta_mut().extend(meta);
|
||
}
|
||
let (responder, receiver) = tokio::sync::oneshot::channel();
|
||
self.tx
|
||
.send(PeerSinkMessage::Request {
|
||
request,
|
||
id: id.clone(),
|
||
responder,
|
||
})
|
||
.await
|
||
.map_err(|_m| ServiceError::TransportClosed)?;
|
||
Ok(RequestHandle {
|
||
id,
|
||
rx: receiver,
|
||
progress_token,
|
||
options,
|
||
peer: self.clone(),
|
||
})
|
||
}
|
||
pub fn peer_info(&self) -> Option<&R::PeerInfo> {
|
||
self.info.get()
|
||
}
|
||
|
||
pub fn set_peer_info(&self, info: R::PeerInfo) {
|
||
if self.info.initialized() {
|
||
tracing::warn!("trying to set peer info, which is already initialized");
|
||
} else {
|
||
let _ = self.info.set(info);
|
||
}
|
||
}
|
||
|
||
pub fn is_transport_closed(&self) -> bool {
|
||
self.tx.is_closed()
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub struct RunningService<R: ServiceRole, S: Service<R>> {
|
||
service: Arc<S>,
|
||
peer: Peer<R>,
|
||
handle: Option<tokio::task::JoinHandle<QuitReason>>,
|
||
cancellation_token: CancellationToken,
|
||
dg: DropGuard,
|
||
}
|
||
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()
|
||
}
|
||
#[inline]
|
||
pub fn cancellation_token(&self) -> RunningServiceCancellationToken {
|
||
RunningServiceCancellationToken(self.cancellation_token.clone())
|
||
}
|
||
|
||
/// Returns true if the service has been closed or cancelled.
|
||
#[inline]
|
||
pub fn is_closed(&self) -> bool {
|
||
self.handle.is_none() || self.cancellation_token.is_cancelled()
|
||
}
|
||
|
||
/// Wait for the service to complete.
|
||
///
|
||
/// This will block until the service loop terminates (either due to
|
||
/// cancellation, transport closure, or an error).
|
||
#[inline]
|
||
pub async fn waiting(mut self) -> Result<QuitReason, tokio::task::JoinError> {
|
||
match self.handle.take() {
|
||
Some(handle) => handle.await,
|
||
None => Ok(QuitReason::Closed),
|
||
}
|
||
}
|
||
|
||
/// Gracefully close the connection and wait for cleanup to complete.
|
||
///
|
||
/// This method cancels the service, waits for the background task to finish
|
||
/// (which includes calling `transport.close()`), and ensures all cleanup
|
||
/// operations complete before returning.
|
||
///
|
||
/// Unlike [`cancel`](Self::cancel), this method takes `&mut self` and can be
|
||
/// called without consuming the `RunningService`. After calling this method,
|
||
/// the service is considered closed and subsequent operations will fail.
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust,ignore
|
||
/// let mut client = ().serve(transport).await?;
|
||
/// // ... use the client ...
|
||
/// client.close().await?;
|
||
/// ```
|
||
pub async fn close(&mut self) -> Result<QuitReason, tokio::task::JoinError> {
|
||
if let Some(handle) = self.handle.take() {
|
||
// Disarm the drop guard so it doesn't try to cancel again
|
||
// We need to cancel manually and wait for completion
|
||
self.cancellation_token.cancel();
|
||
handle.await
|
||
} else {
|
||
// Already closed
|
||
Ok(QuitReason::Closed)
|
||
}
|
||
}
|
||
|
||
/// Gracefully close the connection with a timeout.
|
||
///
|
||
/// Similar to [`close`](Self::close), but returns after the specified timeout
|
||
/// if the cleanup doesn't complete in time. This is useful for ensuring
|
||
/// a bounded shutdown time.
|
||
///
|
||
/// Returns `Ok(Some(reason))` if shutdown completed within the timeout,
|
||
/// `Ok(None)` if the timeout was reached, or `Err` if there was a join error.
|
||
pub async fn close_with_timeout(
|
||
&mut self,
|
||
timeout: Duration,
|
||
) -> Result<Option<QuitReason>, tokio::task::JoinError> {
|
||
if let Some(handle) = self.handle.take() {
|
||
self.cancellation_token.cancel();
|
||
match tokio::time::timeout(timeout, handle).await {
|
||
Ok(result) => result.map(Some),
|
||
Err(_elapsed) => {
|
||
tracing::warn!(
|
||
"close_with_timeout: cleanup did not complete within {:?}",
|
||
timeout
|
||
);
|
||
Ok(None)
|
||
}
|
||
}
|
||
} else {
|
||
Ok(Some(QuitReason::Closed))
|
||
}
|
||
}
|
||
|
||
/// Cancel the service and wait for cleanup to complete.
|
||
///
|
||
/// This consumes the `RunningService` and ensures the connection is properly
|
||
/// closed. For a non-consuming alternative, see [`close`](Self::close).
|
||
pub async fn cancel(mut self) -> Result<QuitReason, tokio::task::JoinError> {
|
||
// Disarm the drop guard since we're handling cancellation explicitly
|
||
let _ = std::mem::replace(&mut self.dg, self.cancellation_token.clone().drop_guard());
|
||
self.close().await
|
||
}
|
||
}
|
||
|
||
impl<R: ServiceRole, S: Service<R>> Drop for RunningService<R, S> {
|
||
fn drop(&mut self) {
|
||
if self.handle.is_some() && !self.cancellation_token.is_cancelled() {
|
||
tracing::debug!(
|
||
"RunningService dropped without explicit close(). \
|
||
The connection will be closed asynchronously. \
|
||
For guaranteed cleanup, call close() or cancel() before dropping."
|
||
);
|
||
}
|
||
// The DropGuard will handle cancellation
|
||
}
|
||
}
|
||
|
||
// use a wrapper type so we can tweak the implementation if needed
|
||
pub struct RunningServiceCancellationToken(CancellationToken);
|
||
|
||
impl RunningServiceCancellationToken {
|
||
pub fn cancel(self) {
|
||
self.0.cancel();
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
#[non_exhaustive]
|
||
pub enum QuitReason {
|
||
Cancelled,
|
||
Closed,
|
||
JoinError(tokio::task::JoinError),
|
||
}
|
||
|
||
/// 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,
|
||
pub meta: Meta,
|
||
pub extensions: Extensions,
|
||
/// An interface to fetch the remote client or server
|
||
pub peer: Peer<R>,
|
||
}
|
||
|
||
impl<R: ServiceRole> RequestContext<R> {
|
||
/// Create a new RequestContext.
|
||
pub fn new(id: RequestId, peer: Peer<R>) -> Self {
|
||
Self {
|
||
ct: CancellationToken::new(),
|
||
id,
|
||
meta: Meta::default(),
|
||
extensions: Extensions::default(),
|
||
peer,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Request execution context
|
||
#[derive(Debug, Clone)]
|
||
pub struct NotificationContext<R: ServiceRole> {
|
||
pub meta: Meta,
|
||
pub extensions: Extensions,
|
||
/// An interface to fetch the remote client or server
|
||
pub peer: Peer<R>,
|
||
}
|
||
|
||
/// Use this function to skip initialization process
|
||
pub fn serve_directly<R, S, T, E, A>(
|
||
service: S,
|
||
transport: T,
|
||
peer_info: Option<R::PeerInfo>,
|
||
) -> RunningService<R, S>
|
||
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())
|
||
}
|
||
|
||
/// Use this function to skip initialization process
|
||
pub fn serve_directly_with_ct<R, S, T, E, A>(
|
||
service: S,
|
||
transport: T,
|
||
peer_info: Option<R::PeerInfo>,
|
||
ct: CancellationToken,
|
||
) -> RunningService<R, S>
|
||
where
|
||
R: ServiceRole,
|
||
S: Service<R>,
|
||
T: IntoTransport<R, E, A>,
|
||
E: std::error::Error + Send + Sync + 'static,
|
||
{
|
||
let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
|
||
serve_inner(service, transport.into_transport(), peer, peer_rx, ct)
|
||
}
|
||
|
||
/// Spawn a task that may hold `!Send` state when the `local` feature is active.
|
||
///
|
||
/// Without the `local` feature this is `tokio::spawn` (requires `Future: Send + 'static`).
|
||
/// With `local` it uses `tokio::task::spawn_local` (requires only `Future: 'static`).
|
||
#[cfg(not(feature = "local"))]
|
||
fn spawn_service_task<F>(future: F) -> tokio::task::JoinHandle<F::Output>
|
||
where
|
||
F: Future + Send + 'static,
|
||
F::Output: Send + 'static,
|
||
{
|
||
tokio::spawn(future)
|
||
}
|
||
|
||
#[cfg(feature = "local")]
|
||
fn spawn_service_task<F>(future: F) -> tokio::task::JoinHandle<F::Output>
|
||
where
|
||
F: Future + 'static,
|
||
F::Output: 'static,
|
||
{
|
||
tokio::task::spawn_local(future)
|
||
}
|
||
|
||
#[instrument(skip_all)]
|
||
fn serve_inner<R, S, T>(
|
||
service: S,
|
||
transport: T,
|
||
peer: Peer<R>,
|
||
mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
|
||
ct: CancellationToken,
|
||
) -> RunningService<R, S>
|
||
where
|
||
R: ServiceRole,
|
||
S: Service<R>,
|
||
T: Transport<R> + 'static,
|
||
{
|
||
const SINK_PROXY_BUFFER_SIZE: usize = 64;
|
||
let (sink_proxy_tx, mut sink_proxy_rx) =
|
||
tokio::sync::mpsc::channel::<TxJsonRpcMessage<R>>(SINK_PROXY_BUFFER_SIZE);
|
||
let peer_info = peer.peer_info();
|
||
if R::IS_CLIENT {
|
||
tracing::info!(?peer_info, "Service initialized as client");
|
||
} else {
|
||
tracing::info!(?peer_info, "Service initialized as server");
|
||
}
|
||
|
||
let mut local_responder_pool =
|
||
HashMap::<RequestId, Responder<Result<R::PeerResp, ServiceError>>>::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 current_span = tracing::Span::current();
|
||
let handle = spawn_service_task(async move {
|
||
let mut transport = transport.into_transport();
|
||
let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new();
|
||
let mut send_task_set = tokio::task::JoinSet::<SendTaskResult>::new();
|
||
#[derive(Debug)]
|
||
enum SendTaskResult {
|
||
Request {
|
||
id: RequestId,
|
||
result: Result<(), DynamicTransportError>,
|
||
},
|
||
Notification {
|
||
responder: Responder<Result<(), ServiceError>>,
|
||
cancellation_param: Option<CancelledNotificationParam>,
|
||
result: Result<(), DynamicTransportError>,
|
||
},
|
||
}
|
||
#[derive(Debug)]
|
||
enum Event<R: ServiceRole> {
|
||
ProxyMessage(PeerSinkMessage<R>),
|
||
PeerMessage(RxJsonRpcMessage<R>),
|
||
ToSink(TxJsonRpcMessage<R>),
|
||
SendTaskResult(SendTaskResult),
|
||
}
|
||
|
||
let quit_reason = loop {
|
||
let evt = if let Some(m) = batch_messages.pop_front() {
|
||
Event::PeerMessage(m)
|
||
} else {
|
||
tokio::select! {
|
||
m = sink_proxy_rx.recv(), if !sink_proxy_rx.is_closed() => {
|
||
if let Some(m) = m {
|
||
Event::ToSink(m)
|
||
} else {
|
||
continue
|
||
}
|
||
}
|
||
m = transport.receive() => {
|
||
if let Some(m) = m {
|
||
Event::PeerMessage(m)
|
||
} else {
|
||
// input stream closed
|
||
tracing::info!("input stream terminated");
|
||
break QuitReason::Closed
|
||
}
|
||
}
|
||
m = peer_rx.recv(), if !peer_rx.is_closed() => {
|
||
if let Some(m) = m {
|
||
Event::ProxyMessage(m)
|
||
} else {
|
||
continue
|
||
}
|
||
}
|
||
m = send_task_set.join_next(), if !send_task_set.is_empty() => {
|
||
let Some(result) = m else {
|
||
continue
|
||
};
|
||
match result {
|
||
Err(e) => {
|
||
// join error, which is serious, we should quit.
|
||
tracing::error!(%e, "send request task encounter a tokio join error");
|
||
break QuitReason::JoinError(e)
|
||
}
|
||
Ok(result) => {
|
||
Event::SendTaskResult(result)
|
||
}
|
||
}
|
||
}
|
||
_ = serve_loop_ct.cancelled() => {
|
||
tracing::info!("task cancelled");
|
||
break QuitReason::Cancelled
|
||
}
|
||
}
|
||
};
|
||
|
||
tracing::trace!(?evt, "new event");
|
||
match evt {
|
||
Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
|
||
if let Err(e) = result {
|
||
if let Some(responder) = local_responder_pool.remove(&id) {
|
||
let _ = responder.send(Err(ServiceError::TransportSend(e)));
|
||
}
|
||
}
|
||
}
|
||
Event::SendTaskResult(SendTaskResult::Notification {
|
||
responder,
|
||
result,
|
||
cancellation_param,
|
||
}) => {
|
||
let response = if let Err(e) = result {
|
||
Err(ServiceError::TransportSend(e))
|
||
} else {
|
||
Ok(())
|
||
};
|
||
let _ = responder.send(response);
|
||
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(),
|
||
}));
|
||
}
|
||
}
|
||
}
|
||
// response and error
|
||
Event::ToSink(m) => {
|
||
if let Some(id) = match &m {
|
||
JsonRpcMessage::Response(response) => Some(&response.id),
|
||
JsonRpcMessage::Error(error) => Some(&error.id),
|
||
_ => None,
|
||
} {
|
||
if let Some(ct) = local_ct_pool.remove(id) {
|
||
ct.cancel();
|
||
}
|
||
let send = transport.send(m);
|
||
let current_span = tracing::Span::current();
|
||
tokio::spawn(async move {
|
||
let send_result = send.await;
|
||
if let Err(error) = send_result {
|
||
tracing::error!(%error, "fail to response message");
|
||
}
|
||
}.instrument(current_span));
|
||
}
|
||
}
|
||
Event::ProxyMessage(PeerSinkMessage::Request {
|
||
request,
|
||
id,
|
||
responder,
|
||
}) => {
|
||
local_responder_pool.insert(id.clone(), responder);
|
||
let send = transport.send(JsonRpcMessage::request(request, id.clone()));
|
||
{
|
||
let id = id.clone();
|
||
let current_span = tracing::Span::current();
|
||
send_task_set.spawn(send.map(move |r| SendTaskResult::Request {
|
||
id,
|
||
result: r.map_err(DynamicTransportError::new::<T, R>),
|
||
}).instrument(current_span));
|
||
}
|
||
}
|
||
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 = transport.send(JsonRpcMessage::notification(notification));
|
||
let current_span = tracing::Span::current();
|
||
send_task_set.spawn(send.map(move |result| SendTaskResult::Notification {
|
||
responder,
|
||
cancellation_param,
|
||
result: result.map_err(DynamicTransportError::new::<T, R>),
|
||
}).instrument(current_span));
|
||
}
|
||
Event::PeerMessage(JsonRpcMessage::Request(JsonRpcRequest {
|
||
id,
|
||
mut request,
|
||
..
|
||
})) => {
|
||
tracing::debug!(%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 mut extensions = Extensions::new();
|
||
let mut meta = Meta::new();
|
||
// avoid clone
|
||
// swap meta firstly, otherwise progress token will be lost
|
||
std::mem::swap(&mut meta, request.get_meta_mut());
|
||
std::mem::swap(&mut extensions, request.extensions_mut());
|
||
let context = RequestContext {
|
||
ct: context_ct,
|
||
id: id.clone(),
|
||
peer: peer.clone(),
|
||
meta,
|
||
extensions,
|
||
};
|
||
let current_span = tracing::Span::current();
|
||
spawn_service_task(async move {
|
||
let result = service
|
||
.handle_request(request, context)
|
||
.await;
|
||
let response = match result {
|
||
Ok(result) => {
|
||
tracing::debug!(%id, ?result, "response message");
|
||
JsonRpcMessage::response(result, id)
|
||
}
|
||
Err(error) => {
|
||
tracing::warn!(%id, ?error, "response error");
|
||
JsonRpcMessage::error(error, id)
|
||
}
|
||
};
|
||
let _send_result = sink.send(response).await;
|
||
}.instrument(current_span));
|
||
}
|
||
}
|
||
Event::PeerMessage(JsonRpcMessage::Notification(JsonRpcNotification {
|
||
notification,
|
||
..
|
||
})) => {
|
||
tracing::info!(?notification, "received notification");
|
||
// catch cancelled notification
|
||
let mut 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();
|
||
let mut extensions = Extensions::new();
|
||
let mut meta = Meta::new();
|
||
// avoid clone
|
||
std::mem::swap(&mut extensions, notification.extensions_mut());
|
||
std::mem::swap(&mut meta, notification.get_meta_mut());
|
||
let context = NotificationContext {
|
||
peer: peer.clone(),
|
||
meta,
|
||
extensions,
|
||
};
|
||
let current_span = tracing::Span::current();
|
||
spawn_service_task(async move {
|
||
let result = service.handle_notification(notification, context).await;
|
||
if let Err(error) = result {
|
||
tracing::warn!(%error, "Error sending notification");
|
||
}
|
||
}.instrument(current_span));
|
||
}
|
||
}
|
||
Event::PeerMessage(JsonRpcMessage::Response(JsonRpcResponse {
|
||
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(JsonRpcMessage::Error(JsonRpcError { 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 = transport.close().await;
|
||
if let Err(e) = sink_close_result {
|
||
tracing::error!(%e, "fail to close sink");
|
||
}
|
||
tracing::info!(?quit_reason, "serve finished");
|
||
quit_reason
|
||
}.instrument(current_span));
|
||
RunningService {
|
||
service,
|
||
peer: peer_return,
|
||
handle: Some(handle),
|
||
cancellation_token: ct.clone(),
|
||
dg: ct.drop_guard(),
|
||
}
|
||
}
|