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

feat: throw initialize error detail (#192)

Browse source
This commit is contained in:
4t145 2025-05-19 09:41:56 +08:00 • committed by GitHub
commit 817e62d895
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 115 additions and 93 deletions
Download patch file Download diff file Expand all files Collapse all files

7
crates/rmcp/src/model.rs
View file

@ -119,6 +119,13 @@ impl Default for ProtocolVersion {
Self::LATEST
}
}
impl std::fmt::Display for ProtocolVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.0.fmt(f)
}
}
impl ProtocolVersion {
pub const V_2025_03_26: Self = Self(Cow::Borrowed("2025-03-26"));
pub const V_2024_11_05: Self = Self(Cow::Borrowed("2024-11-05"));

15
crates/rmcp/src/service.rs
View file

@ -77,6 +77,7 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
+ From<CancelledNotification>
+ TransferObject;
type InitializeError<E>;
const IS_CLIENT: bool;
type Info: TransferObject;
type PeerInfo: TransferObject;
@ -113,7 +114,7 @@ pub trait ServiceExt<R: ServiceRole>: Service<R> + Sized {
fn serve<T, E, A>(
self,
transport: T,
) -> impl Future<Output = Result<RunningService<R, Self>, E>> + Send
) -> impl Future<Output = Result<RunningService<R, Self>, R::InitializeError<E>>> + Send
where
T: IntoTransport<R, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
@ -125,7 +126,7 @@ pub trait ServiceExt<R: ServiceRole>: Service<R> + Sized {
self,
transport: T,
ct: CancellationToken,
) -> impl Future<Output = Result<RunningService<R, Self>, E>> + Send
) -> impl Future<Output = Result<RunningService<R, Self>, R::InitializeError<E>>> + Send
where
T: IntoTransport<R, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
@ -469,7 +470,7 @@ pub async fn serve_directly<R, S, T, E, A>(
service: S,
transport: T,
peer_info: R::PeerInfo,
) -> Result<RunningService<R, S>, E>
) -> RunningService<R, S>
where
R: ServiceRole,
S: Service<R>,
@ -485,7 +486,7 @@ pub async fn serve_directly_with_ct<R, S, T, E, A>(
transport: T,
peer_info: R::PeerInfo,
ct: CancellationToken,
) -> Result<RunningService<R, S>, E>
) -> RunningService<R, S>
where
R: ServiceRole,
S: Service<R>,
@ -503,7 +504,7 @@ async fn serve_inner<R, S, T, E, A>(
peer: Peer<R>,
mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
ct: CancellationToken,
) -> Result<RunningService<R, S>, E>
) -> RunningService<R, S>
where
R: ServiceRole,
S: Service<R>,
@ -788,10 +789,10 @@ where
tracing::info!(?quit_reason, "serve finished");
quit_reason
});
Ok(RunningService {
RunningService {
service,
peer: peer_return,
handle,
dg: ct.drop_guard(),
})
}
}

68
crates/rmcp/src/service/client.rs
View file

@ -1,3 +1,5 @@
use std::borrow::Cow;
use thiserror::Error;
use super::*;
@ -19,7 +21,7 @@ use crate::model::{
///
/// if you want to handle the error, you can use `serve_client_with_ct` or `serve_client` with `Result<RunningService<RoleClient, S>, ClientError>`
#[derive(Error, Debug)]
pub enum ClientError {
pub enum ClientInitializeError<E> {
#[error("expect initialized response, but received: {0:?}")]
ExpectedInitResponse(Option<ServerJsonRpcMessage>),
@ -32,30 +34,32 @@ pub enum ClientError {
#[error("connection closed: {0}")]
ConnectionClosed(String),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Send message error {error}, when {context}")]
TransportError {
error: E,
context: Cow<'static, str>,
},
}
/// Helper function to get the next message from the stream
async fn expect_next_message<T>(
async fn expect_next_message<T, E>(
transport: &mut T,
context: &str,
) -> Result<ServerJsonRpcMessage, ClientError>
) -> Result<ServerJsonRpcMessage, ClientInitializeError<E>>
where
T: Transport<RoleClient>,
{
transport
.receive()
.await
.ok_or_else(|| ClientError::ConnectionClosed(context.to_string()))
.map_err(|e| ClientError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))
.ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string()))
}
/// Helper function to expect a response from the stream
async fn expect_response<T>(
async fn expect_response<T, E>(
transport: &mut T,
context: &str,
) -> Result<(ServerResult, RequestId), ClientError>
) -> Result<(ServerResult, RequestId), ClientInitializeError<E>>
where
T: Transport<RoleClient>,
{
@ -63,7 +67,7 @@ where
match msg {
ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => Ok((result, id)),
_ => Err(ClientError::ExpectedInitResponse(Some(msg))),
_ => Err(ClientInitializeError::ExpectedInitResponse(Some(msg))),
}
}
@ -79,7 +83,7 @@ impl ServiceRole for RoleClient {
type PeerNot = ServerNotification;
type Info = ClientInfo;
type PeerInfo = ServerInfo;
type InitializeError<E> = ClientInitializeError<E>;
const IS_CLIENT: bool = true;
}
@ -90,7 +94,7 @@ impl<S: Service<RoleClient>> ServiceExt<RoleClient> for S {
self,
transport: T,
ct: CancellationToken,
) -> impl Future<Output = Result<RunningService<RoleClient, Self>, E>> + Send
) -> impl Future<Output = Result<RunningService<RoleClient, Self>, ClientInitializeError<E>>> + Send
where
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
@ -103,7 +107,7 @@ impl<S: Service<RoleClient>> ServiceExt<RoleClient> for S {
pub async fn serve_client<S, T, E, A>(
service: S,
transport: T,
) -> Result<RunningService<RoleClient, S>, E>
) -> Result<RunningService<RoleClient, S>, ClientInitializeError<E>>
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
@ -116,7 +120,7 @@ pub async fn serve_client_with_ct<S, T, E, A>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleClient, S>, E>
) -> Result<RunningService<RoleClient, S>, ClientInitializeError<E>>
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
@ -125,14 +129,6 @@ where
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// Convert ClientError to std::io::Error, then to E
let handle_client_error = |e: ClientError| -> E {
match e {
ClientError::Io(io_err) => io_err.into(),
other => std::io::Error::new(std::io::ErrorKind::Other, format!("{}", other)).into(),
}
};
// service
let id = id_provider.next_request_id();
let init_request = InitializeRequest {
@ -145,23 +141,23 @@ where
ClientRequest::InitializeRequest(init_request),
id.clone(),
))
.await?;
let (response, response_id) = expect_response(&mut transport, "initialize response")
.await
.map_err(handle_client_error)?;
.map_err(|error| ClientInitializeError::TransportError {
error,
context: "send initialize request".into(),
})?;
let (response, response_id) = expect_response(&mut transport, "initialize response").await?;
if id != response_id {
return Err(handle_client_error(ClientError::ConflictInitResponseId(
return Err(ClientInitializeError::ConflictInitResponseId(
id,
response_id,
)));
));
}
let ServerResult::InitializeResult(initialize_result) = response else {
return Err(handle_client_error(ClientError::ExpectedInitResult(Some(
response,
))));
return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
};
// send notification
@ -171,9 +167,15 @@ where
extensions: Default::default(),
}),
);
transport.send(notification).await?;
transport
.send(notification)
.await
.map_err(|error| ClientInitializeError::TransportError {
error,
context: "send initialized notification".into(),
})?;
let (peer, peer_rx) = Peer::new(id_provider, initialize_result);
serve_inner(service, transport, peer, peer_rx, ct).await
Ok(serve_inner(service, transport, peer, peer_rx, ct).await)
}
macro_rules! method {

108
crates/rmcp/src/service/server.rs
View file

@ -1,3 +1,5 @@
use std::borrow::Cow;
use thiserror::Error;
use super::*;
@ -6,9 +8,9 @@ use crate::model::{
ClientNotification, ClientRequest, ClientResult, CreateMessageRequest,
CreateMessageRequestParam, CreateMessageResult, ErrorData, ListRootsRequest, ListRootsResult,
LoggingMessageNotification, LoggingMessageNotificationParam, ProgressNotification,
ProgressNotificationParam, PromptListChangedNotification, ResourceListChangedNotification,
ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ServerInfo, ServerNotification,
ServerRequest, ServerResult, ToolListChangedNotification,
ProgressNotificationParam, PromptListChangedNotification, ProtocolVersion,
ResourceListChangedNotification, ResourceUpdatedNotification, ResourceUpdatedNotificationParam,
ServerInfo, ServerNotification, ServerRequest, ServerResult, ToolListChangedNotification,
};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
@ -23,6 +25,8 @@ impl ServiceRole for RoleServer {
type PeerNot = ClientNotification;
type Info = ServerInfo;
type PeerInfo = ClientInfo;
type InitializeError<E> = ServerInitializeError<E>;
const IS_CLIENT: bool = false;
}
@ -30,12 +34,12 @@ impl ServiceRole for RoleServer {
///
/// if you want to handle the error, you can use `serve_server_with_ct` or `serve_server` with `Result<RunningService<RoleServer, S>, ServerError>`
#[derive(Error, Debug)]
pub enum ServerError {
pub enum ServerInitializeError<E> {
#[error("expect initialized request, but received: {0:?}")]
ExpectedInitRequest(Option<ClientJsonRpcMessage>),
ExpectedInitializeRequest(Option<ClientJsonRpcMessage>),
#[error("expect initialized notification, but received: {0:?}")]
ExpectedInitNotification(Option<ClientJsonRpcMessage>),
ExpectedInitializedNotification(Option<ClientJsonRpcMessage>),
#[error("connection closed: {0}")]
ConnectionClosed(String),
@ -46,8 +50,14 @@ pub enum ServerError {
#[error("initialize failed: {0}")]
InitializeFailed(ErrorData),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("unsupported protocol version: {0}")]
UnsupportedProtocolVersion(ProtocolVersion),
#[error("Send message error {error}, when {context}")]
TransportError {
error: E,
context: Cow<'static, str>,
},
}
pub type ClientSink = Peer<RoleServer>;
@ -57,7 +67,7 @@ impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
self,
transport: T,
ct: CancellationToken,
) -> impl Future<Output = Result<RunningService<RoleServer, Self>, E>> + Send
) -> impl Future<Output = Result<RunningService<RoleServer, Self>, ServerInitializeError<E>>> + Send
where
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
@ -70,7 +80,7 @@ impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
pub async fn serve_server<S, T, E, A>(
service: S,
transport: T,
) -> Result<RunningService<RoleServer, S>, E>
) -> Result<RunningService<RoleServer, S>, ServerInitializeError<E>>
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
@ -80,52 +90,56 @@ where
}
/// Helper function to get the next message from the stream
async fn expect_next_message<T>(
async fn expect_next_message<T, E>(
transport: &mut T,
context: &str,
) -> Result<ClientJsonRpcMessage, ServerError>
) -> Result<ClientJsonRpcMessage, ServerInitializeError<E>>
where
T: Transport<RoleServer>,
{
transport
.receive()
.await
.ok_or_else(|| ServerError::ConnectionClosed(context.to_string()))
.ok_or_else(|| ServerInitializeError::ConnectionClosed(context.to_string()))
}
/// Helper function to expect a request from the stream
async fn expect_request<T>(
async fn expect_request<T, E>(
transport: &mut T,
context: &str,
) -> Result<(ClientRequest, RequestId), ServerError>
) -> Result<(ClientRequest, RequestId), ServerInitializeError<E>>
where
T: Transport<RoleServer>,
{
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_request()
.ok_or(ServerError::ExpectedInitRequest(Some(msg_clone)))
.ok_or(ServerInitializeError::ExpectedInitializeRequest(Some(
msg_clone,
)))
}
/// Helper function to expect a notification from the stream
async fn expect_notification<T>(
async fn expect_notification<T, E>(
transport: &mut T,
context: &str,
) -> Result<ClientNotification, ServerError>
) -> Result<ClientNotification, ServerInitializeError<E>>
where
T: Transport<RoleServer>,
{
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_notification()
.ok_or(ServerError::ExpectedInitNotification(Some(msg_clone)))
.ok_or(ServerInitializeError::ExpectedInitializedNotification(
Some(msg_clone),
))
}
pub async fn serve_server_with_ct<S, T, E, A>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleServer, S>, E>
) -> Result<RunningService<RoleServer, S>, ServerInitializeError<E>>
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
@ -134,23 +148,13 @@ where
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// Convert ServerError to std::io::Error, then to E
let handle_server_error = |e: ServerError| -> E {
match e {
ServerError::Io(io_err) => io_err.into(),
other => std::io::Error::new(std::io::ErrorKind::Other, format!("{}", other)).into(),
}
};
// Get initialize request
let (request, id) = expect_request(&mut transport, "initialized request")
.await
.map_err(handle_server_error)?;
let (request, id) = expect_request(&mut transport, "initialized request").await?;
let ClientRequest::InitializeRequest(peer_info) = &request else {
return Err(handle_server_error(ServerError::ExpectedInitRequest(Some(
return Err(ServerInitializeError::ExpectedInitializeRequest(Some(
ClientJsonRpcMessage::request(request, id),
))));
)));
};
let (peer, peer_rx) = Peer::new(id_provider, peer_info.params.clone());
let context = RequestContext {
@ -165,24 +169,24 @@ where
let mut init_response = match init_response {
Ok(ServerResult::InitializeResult(init_response)) => init_response,
Ok(result) => {
return Err(handle_server_error(
ServerError::UnexpectedInitializeResponse(result),
));
return Err(ServerInitializeError::UnexpectedInitializeResponse(result));
}
Err(e) => {
transport
.send(ServerJsonRpcMessage::error(e.clone(), id))
.await?;
return Err(handle_server_error(ServerError::InitializeFailed(e)));
.await
.map_err(|error| ServerInitializeError::TransportError {
error,
context: "sending error response".into(),
})?;
return Err(ServerInitializeError::InitializeFailed(e));
}
};
let protocol_version = match peer_info
.params
.protocol_version
let peer_protocol_version = peer_info.params.protocol_version.clone();
let protocol_version = match peer_protocol_version
.partial_cmp(&init_response.protocol_version)
.ok_or(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"unsupported protocol version",
.ok_or(ServerInitializeError::UnsupportedProtocolVersion(
peer_protocol_version,
))? {
std::cmp::Ordering::Less => peer_info.params.protocol_version.clone(),
_ => init_response.protocol_version,
@ -193,20 +197,22 @@ where
ServerResult::InitializeResult(init_response),
id,
))
.await?;
.await
.map_err(|error| ServerInitializeError::TransportError {
error,
context: "sending initialize response".into(),
})?;
// Wait for initialize notification
let notification = expect_notification(&mut transport, "initialize notification")
.await
.map_err(handle_server_error)?;
let notification = expect_notification(&mut transport, "initialize notification").await?;
let ClientNotification::InitializedNotification(_) = notification else {
return Err(handle_server_error(ServerError::ExpectedInitNotification(
return Err(ServerInitializeError::ExpectedInitializedNotification(
Some(ClientJsonRpcMessage::notification(notification)),
)));
));
};
let _ = service.handle_notification(notification).await;
// Continue processing service
serve_inner(service, transport, peer, peer_rx, ct).await
Ok(serve_inner(service, transport, peer, peer_rx, ct).await)
}
macro_rules! method {

5
crates/rmcp/src/transport/sse_server.rs
View file

@ -288,7 +288,10 @@ impl SseServer {
let service = service_provider();
let ct = self.config.ct.child_token();
tokio::spawn(async move {
let server = service.serve_with_ct(transport, ct).await?;
let server = service
.serve_with_ct(transport, ct)
.await
.map_err(std::io::Error::other)?;
server.waiting().await?;
tokio::io::Result::Ok(())
});

5
crates/rmcp/src/transport/streamable_http_server/axum.rs
View file

@ -327,7 +327,10 @@ impl StreamableHttpServer {
let service = service_provider();
let ct = self.config.ct.child_token();
tokio::spawn(async move {
let server = service.serve_with_ct(transport, ct).await?;
let server = service
.serve_with_ct(transport, ct)
.await
.map_err(tokio::io::Error::other)?;
server.waiting().await?;
tokio::io::Result::Ok(())
});