modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

Move whole rmcp crate to offcial rust sdk (#44)

Browse source 
* get rid of async-trait

* move rmcp to official-rust-sdk

* replace master with main

* update readme, change license

* fix typo

* Update README.md

Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>

* fix typos

---------

Co-authored-by: = <=>
Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>
This commit is contained in:
4t145 2025-03-28 05:33:41 +08:00 • committed by GitHub
commit 0be6067326
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 8676 additions and 4466 deletions
Download patch file Download diff file Expand all files Collapse all files

287
crates/rmcp/src/service/client.rs Normal file
View file

@ -0,0 +1,287 @@
use crate::model::{
CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
CancelledNotificationParam, ClientInfo, ClientMessage, ClientNotification, ClientRequest,
ClientResult, CompleteRequest, CompleteRequestParam, CompleteResult, GetPromptRequest,
GetPromptRequestParam, GetPromptResult, InitializeRequest, InitializedNotification,
ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest,
ListToolsResult, PaginatedRequestParam, PaginatedRequestParamInner, ProgressNotification,
ProgressNotificationParam, ReadResourceRequest, ReadResourceRequestParam, ReadResourceResult,
RootsListChangedNotification, ServerInfo, ServerNotification, ServerRequest, ServerResult,
SetLevelRequest, SetLevelRequestParam, SubscribeRequest, SubscribeRequestParam,
UnsubscribeRequest, UnsubscribeRequestParam,
};
use super::*;
use futures::{SinkExt, StreamExt};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RoleClient;
impl ServiceRole for RoleClient {
type Req = ClientRequest;
type Resp = ClientResult;
type Not = ClientNotification;
type PeerReq = ServerRequest;
type PeerResp = ServerResult;
type PeerNot = ServerNotification;
type Info = ClientInfo;
type PeerInfo = ServerInfo;
const IS_CLIENT: bool = true;
}
pub type ServerSink = Peer<RoleClient>;
impl<S: Service<RoleClient>> ServiceExt<RoleClient> for S {
fn serve_with_ct<T, E, A>(
self,
transport: T,
ct: CancellationToken,
) -> impl Future<Output = Result<RunningService<RoleClient, Self>, E>> + Send
where
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
Self: Sized,
{
serve_client_with_ct(self, transport, ct)
}
}
pub async fn serve_client<S, T, E, A>(
service: S,
transport: T,
) -> Result<RunningService<RoleClient, S>, E>
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
{
serve_client_with_ct(service, transport, Default::default()).await
}
pub async fn serve_client_with_ct<S, T, E, A>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleClient, S>, E>
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
{
let (sink, stream) = transport.into_transport();
let mut sink = Box::pin(sink);
let mut stream = Box::pin(stream);
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// service
let id = id_provider.next_request_id();
let init_request = InitializeRequest {
method: Default::default(),
params: service.get_info(),
};
sink.send(
ClientMessage::Request(ClientRequest::InitializeRequest(init_request), id.clone())
.into_json_rpc_message(),
)
.await?;
let (response, response_id) = stream
.next()
.await
.ok_or(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"expect initialize response",
))?
.into_message()
.into_result()
.ok_or(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"expect initialize result",
))?;
if id != response_id {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"conflict initialize response id",
)
.into());
}
let response = response.map_err(std::io::Error::other)?;
let ServerResult::InitializeResult(initialize_result) = response else {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"expect initialize result",
)
.into());
};
// send notification
let notification = ClientMessage::Notification(ClientNotification::InitializedNotification(
InitializedNotification {
method: Default::default(),
},
));
sink.send(notification.into_json_rpc_message()).await?;
serve_inner(service, (sink, stream), initialize_result, id_provider, ct).await
}
macro_rules! method {
(peer_req $method:ident $Req:ident() => $Resp: ident ) => {
pub async fn $method(&self) -> Result<$Resp, ServiceError> {
let result = self
.send_request(ClientRequest::$Req($Req {
method: Default::default(),
}))
.await?;
match result {
ServerResult::$Resp(result) => Ok(result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
(peer_req $method:ident $Req:ident($Param: ident) => $Resp: ident ) => {
pub async fn $method(&self, params: $Param) -> Result<$Resp, ServiceError> {
let result = self
.send_request(ClientRequest::$Req($Req {
method: Default::default(),
params,
}))
.await?;
match result {
ServerResult::$Resp(result) => Ok(result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
(peer_req $method:ident $Req:ident($Param: ident)) => {
pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
let result = self
.send_request(ClientRequest::$Req($Req {
method: Default::default(),
params,
}))
.await?;
match result {
ServerResult::EmptyResult(_) => Ok(()),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
(peer_not $method:ident $Not:ident($Param: ident)) => {
pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
self.send_notification(ClientNotification::$Not($Not {
method: Default::default(),
params,
}))
.await?;
Ok(())
}
};
(peer_not $method:ident $Not:ident) => {
pub async fn $method(&self) -> Result<(), ServiceError> {
self.send_notification(ClientNotification::$Not($Not {
method: Default::default(),
}))
.await?;
Ok(())
}
};
}
impl Peer<RoleClient> {
method!(peer_req complete CompleteRequest(CompleteRequestParam) => CompleteResult);
method!(peer_req set_level SetLevelRequest(SetLevelRequestParam));
method!(peer_req get_prompt GetPromptRequest(GetPromptRequestParam) => GetPromptResult);
method!(peer_req list_prompts ListPromptsRequest(PaginatedRequestParam) => ListPromptsResult);
method!(peer_req list_resources ListResourcesRequest(PaginatedRequestParam) => ListResourcesResult);
method!(peer_req list_resource_templates ListResourceTemplatesRequest(PaginatedRequestParam) => ListResourceTemplatesResult);
method!(peer_req read_resource ReadResourceRequest(ReadResourceRequestParam) => ReadResourceResult);
method!(peer_req subscribe SubscribeRequest(SubscribeRequestParam) );
method!(peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParam));
method!(peer_req call_tool CallToolRequest(CallToolRequestParam) => CallToolResult);
method!(peer_req list_tools ListToolsRequest(PaginatedRequestParam) => ListToolsResult);
method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam));
method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam));
method!(peer_not notify_initialized InitializedNotification);
method!(peer_not notify_roots_list_changed RootsListChangedNotification);
}
impl Peer<RoleClient> {
/// A wrapper method for [`Peer<RoleClient>::list_tools`].
///
/// This function will call [`Peer<RoleClient>::list_tools`] multiple times until all tools are listed.
pub async fn list_all_tools(&self) -> Result<Vec<crate::model::Tool>, ServiceError> {
let mut tools = Vec::new();
let mut cursor = None;
loop {
let result = self
.list_tools(Some(PaginatedRequestParamInner { cursor }))
.await?;
tools.extend(result.tools);
cursor = result.next_cursor;
if cursor.is_none() {
break;
}
}
Ok(tools)
}
/// A wrapper method for [`Peer<RoleClient>::list_resources`].
///
/// This function will call [`Peer<RoleClient>::list_resources`] multiple times until all resources are listed.
pub async fn list_all_prompts(&self) -> Result<Vec<crate::model::Prompt>, ServiceError> {
let mut prompts = Vec::new();
let mut cursor = None;
loop {
let result = self
.list_prompts(Some(PaginatedRequestParamInner { cursor }))
.await?;
prompts.extend(result.prompts);
cursor = result.next_cursor;
if cursor.is_none() {
break;
}
}
Ok(prompts)
}
/// A wrapper method for [`Peer<RoleClient>::list_resources`].
///
/// This function will call [`Peer<RoleClient>::list_resources`] multiple times until all resources are listed.
pub async fn list_all_resources(&self) -> Result<Vec<crate::model::Resource>, ServiceError> {
let mut resources = Vec::new();
let mut cursor = None;
loop {
let result = self
.list_resources(Some(PaginatedRequestParamInner { cursor }))
.await?;
resources.extend(result.resources);
cursor = result.next_cursor;
if cursor.is_none() {
break;
}
}
Ok(resources)
}
/// A wrapper method for [`Peer<RoleClient>::list_resource_templates`].
///
/// This function will call [`Peer<RoleClient>::list_resource_templates`] multiple times until all resource templates are listed.
pub async fn list_all_resource_templates(
&self,
) -> Result<Vec<crate::model::ResourceTemplate>, ServiceError> {
let mut resource_templates = Vec::new();
let mut cursor = None;
loop {
let result = self
.list_resource_templates(Some(PaginatedRequestParamInner { cursor }))
.await?;
resource_templates.extend(result.resource_templates);
cursor = result.next_cursor;
if cursor.is_none() {
break;
}
}
Ok(resource_templates)
}
}

204
crates/rmcp/src/service/server.rs Normal file
View file

@ -0,0 +1,204 @@
use crate::model::{
CancelledNotification, CancelledNotificationParam, ClientInfo, ClientNotification,
ClientRequest, ClientResult, CreateMessageRequest, CreateMessageRequestParam,
CreateMessageResult, ListRootsRequest, ListRootsResult, LoggingMessageNotification,
LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam,
PromptListChangedNotification, ResourceListChangedNotification, ResourceUpdatedNotification,
ResourceUpdatedNotificationParam, ServerInfo, ServerMessage, ServerNotification, ServerRequest,
ServerResult, ToolListChangedNotification,
};
use super::*;
use futures::{SinkExt, StreamExt};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RoleServer;
impl ServiceRole for RoleServer {
type Req = ServerRequest;
type Resp = ServerResult;
type Not = ServerNotification;
type PeerReq = ClientRequest;
type PeerResp = ClientResult;
type PeerNot = ClientNotification;
type Info = ServerInfo;
type PeerInfo = ClientInfo;
const IS_CLIENT: bool = false;
}
pub type ClientSink = Peer<RoleServer>;
impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
fn serve_with_ct<T, E, A>(
self,
transport: T,
ct: CancellationToken,
) -> impl Future<Output = Result<RunningService<RoleServer, Self>, E>> + Send
where
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
Self: Sized,
{
serve_server_with_ct(self, transport, ct)
}
}
pub async fn serve_server<S, T, E, A>(
service: S,
transport: T,
) -> Result<RunningService<RoleServer, S>, E>
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
{
serve_server_with_ct(service, transport, CancellationToken::new()).await
}
pub async fn serve_server_with_ct<S, T, E, A>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleServer, S>, E>
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
{
let (sink, stream) = transport.into_transport();
let mut sink = Box::pin(sink);
let mut stream = Box::pin(stream);
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// service
let (request, id) = stream
.next()
.await
.ok_or(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"expect initialize request",
))?
.into_message()
.into_request()
.ok_or(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"expect initialize request",
))?;
let ClientRequest::InitializeRequest(peer_info) = request else {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"expect initialize request",
)
.into());
};
let init_response = service.get_info();
sink.send(
ServerMessage::Response(ServerResult::InitializeResult(init_response), id)
.into_json_rpc_message(),
)
.await?;
// waiting for notification
let notification = stream
.next()
.await
.ok_or(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"expect initialize notification",
))?
.into_message()
.into_notification()
.ok_or(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"expect initialize notification",
))?;
let ClientNotification::InitializedNotification(_) = notification else {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"expect initialize notification",
)
.into());
};
serve_inner(service, (sink, stream), peer_info.params, id_provider, ct).await
}
macro_rules! method {
(peer_req $method:ident $Req:ident() => $Resp: ident ) => {
pub async fn $method(&self) -> Result<$Resp, ServiceError> {
let result = self
.send_request(ServerRequest::$Req($Req {
method: Default::default(),
}))
.await?;
match result {
ClientResult::$Resp(result) => Ok(result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
(peer_req $method:ident $Req:ident($Param: ident) => $Resp: ident ) => {
pub async fn $method(&self, params: $Param) -> Result<$Resp, ServiceError> {
let result = self
.send_request(ServerRequest::$Req($Req {
method: Default::default(),
params,
}))
.await?;
match result {
ClientResult::$Resp(result) => Ok(result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
(peer_req $method:ident $Req:ident($Param: ident)) => {
pub fn $method(
&self,
params: $Param,
) -> impl Future<Output = Result<(), ServiceError>> + Send + '_ {
async move {
let result = self
.send_request(ServerRequest::$Req($Req {
method: Default::default(),
params,
}))
.await?;
match result {
ClientResult::EmptyResult(_) => Ok(()),
_ => Err(ServiceError::UnexpectedResponse),
}
}
}
};
(peer_not $method:ident $Not:ident($Param: ident)) => {
pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
self.send_notification(ServerNotification::$Not($Not {
method: Default::default(),
params,
}))
.await?;
Ok(())
}
};
(peer_not $method:ident $Not:ident) => {
pub async fn $method(&self) -> Result<(), ServiceError> {
self.send_notification(ServerNotification::$Not($Not {
method: Default::default(),
}))
.await?;
Ok(())
}
};
}
impl Peer<RoleServer> {
method!(peer_req create_message CreateMessageRequest(CreateMessageRequestParam) => CreateMessageResult);
method!(peer_req list_roots ListRootsRequest() => ListRootsResult);
method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam));
method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam));
method!(peer_not notify_logging_message LoggingMessageNotification(LoggingMessageNotificationParam));
method!(peer_not notify_resource_updated ResourceUpdatedNotification(ResourceUpdatedNotificationParam));
method!(peer_not notify_resource_list_changed ResourceListChangedNotification);
method!(peer_not notify_tool_list_changed ToolListChangedNotification);
method!(peer_not notify_prompt_list_changed PromptListChangedNotification);
}

62
crates/rmcp/src/service/tower.rs Normal file
View file

@ -0,0 +1,62 @@
use std::{future::poll_fn, marker::PhantomData};
use tower_service::Service as TowerService;
use crate::service::{Peer, RequestContext, Service, ServiceRole};
pub struct TowerHandler<S, R: ServiceRole> {
pub service: S,
pub info: R::Info,
pub peer: Option<Peer<R>>,
role: PhantomData<R>,
}
impl<S, R: ServiceRole> TowerHandler<S, R> {
pub fn new(service: S, info: R::Info) -> Self {
Self {
service,
role: PhantomData,
info,
peer: None,
}
}
}
impl<S, R: ServiceRole> Service<R> for TowerHandler<S, R>
where
S: TowerService<R::PeerReq, Response = R::Resp> + Sync + Send + Clone + 'static,
S::Error: Into<crate::Error>,
S::Future: Send,
{
async fn handle_request(
&self,
request: R::PeerReq,
_context: RequestContext<R>,
) -> Result<R::Resp, crate::Error> {
let mut service = self.service.clone();
poll_fn(|cx| service.poll_ready(cx))
.await
.map_err(Into::into)?;
let resp = service.call(request).await.map_err(Into::into)?;
Ok(resp)
}
fn handle_notification(
&self,
_notification: R::PeerNot,
) -> impl Future<Output = Result<(), crate::Error>> + Send + '_ {
std::future::ready(Ok(()))
}
fn get_peer(&self) -> Option<Peer<R>> {
self.peer.clone()
}
fn set_peer(&mut self, peer: Peer<R>) {
self.peer = Some(peer);
}
fn get_info(&self) -> R::Info {
self.info.clone()
}
}