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: keep internal error in worker's quit reason (#372)

Browse source 
* feat: support downcast WorkerQuitReason::Fatal

* feat(transport): expose internal worker error in fatal

* Update crates/rmcp/src/transport/streamable_http_client.rs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update crates/rmcp/src/transport/streamable_http_client.rs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
4t145 2025-08-18 11:13:10 +08:00 • committed by GitHub
commit 0a2f77de41
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 71 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

2
crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs
View file

@ -36,7 +36,7 @@ impl StreamableHttpClient for reqwest::Client {
} }
let response = request_builder.send().await?; let response = request_builder.send().await?;
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED { if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
return Err(StreamableHttpError::SeverDoesNotSupportSse); return Err(StreamableHttpError::ServerDoesNotSupportSse);
} }
let response = response.error_for_status()?; let response = response.error_for_status()?;
match response.headers().get(reqwest::header::CONTENT_TYPE) { match response.headers().get(reqwest::header::CONTENT_TYPE) {

25
crates/rmcp/src/transport/streamable_http_client.rs
View file

@ -33,15 +33,17 @@ pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
#[error("Unexpected content type: {0:?}")] #[error("Unexpected content type: {0:?}")]
UnexpectedContentType(Option<String>), UnexpectedContentType(Option<String>),
#[error("Server does not support SSE")] #[error("Server does not support SSE")]
SeverDoesNotSupportSse, ServerDoesNotSupportSse,
#[error("Server does not support delete session")] #[error("Server does not support delete session")]
SeverDoesNotSupportDeleteSession, ServerDoesNotSupportDeleteSession,
#[error("Tokio join error: {0}")] #[error("Tokio join error: {0}")]
TokioJoinError(#[from] tokio::task::JoinError), TokioJoinError(#[from] tokio::task::JoinError),
#[error("Deserialize error: {0}")] #[error("Deserialize error: {0}")]
Deserialize(#[from] serde_json::Error), Deserialize(#[from] serde_json::Error),
#[error("Transport channel closed")] #[error("Transport channel closed")]
TransportChannelClosed, TransportChannelClosed,
#[error("Missing session id in HTTP response")]
MissingSessionIdInResponse,
#[cfg(feature = "auth")] #[cfg(feature = "auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))] #[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
#[error("Auth error: {0}")] #[error("Auth error: {0}")]
@ -54,6 +56,11 @@ impl From<reqwest::Error> for StreamableHttpError<reqwest::Error> {
} }
} }
#[derive(Debug, Clone, Error)]
pub enum StreamableHttpProtocolError {
#[error("Missing session id in response")]
MissingSessionIdInResponse,
}
pub enum StreamableHttpPostResponse { pub enum StreamableHttpPostResponse {
Accepted, Accepted,
Json(ServerJsonRpcMessage, Option<String>), Json(ServerJsonRpcMessage, Option<String>),
@ -261,7 +268,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
async fn run( async fn run(
self, self,
mut context: super::worker::WorkerContext<Self>, mut context: super::worker::WorkerContext<Self>,
) -> Result<(), WorkerQuitReason> { ) -> Result<(), WorkerQuitReason<Self::Error>> {
let channel_buffer_capacity = self.config.channel_buffer_capacity; let channel_buffer_capacity = self.config.channel_buffer_capacity;
let (sse_worker_tx, mut sse_worker_rx) = let (sse_worker_tx, mut sse_worker_rx) =
tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(channel_buffer_capacity); tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(channel_buffer_capacity);
@ -278,7 +285,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
.post_message(config.uri.clone(), initialize_request, None, None) .post_message(config.uri.clone(), initialize_request, None, None)
.await .await
.map_err(WorkerQuitReason::fatal_context("send initialize request"))? .map_err(WorkerQuitReason::fatal_context("send initialize request"))?
.expect_initialized::<Self::Error>() .expect_initialized::<C::Error>()
.await .await
.map_err(WorkerQuitReason::fatal_context( .map_err(WorkerQuitReason::fatal_context(
"process initialize response", "process initialize response",
@ -288,7 +295,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
} else { } else {
if !self.config.allow_stateless { if !self.config.allow_stateless {
return Err(WorkerQuitReason::fatal( return Err(WorkerQuitReason::fatal(
"missing session id in initialize response", StreamableHttpError::<C::Error>::MissingSessionIdInResponse,
"process initialize response", "process initialize response",
)); ));
} }
@ -308,7 +315,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
Ok(_) => { Ok(_) => {
tracing::info!(session_id = session_id.as_ref(), "delete session success") tracing::info!(session_id = session_id.as_ref(), "delete session success")
} }
Err(StreamableHttpError::SeverDoesNotSupportDeleteSession) => { Err(StreamableHttpError::ServerDoesNotSupportDeleteSession) => {
tracing::info!( tracing::info!(
session_id = session_id.as_ref(), session_id = session_id.as_ref(),
"server doesn't support delete session" "server doesn't support delete session"
@ -338,7 +345,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
.map_err(WorkerQuitReason::fatal_context( .map_err(WorkerQuitReason::fatal_context(
"send initialized notification", "send initialized notification",
))? ))?
.expect_accepted::<Self::Error>() .expect_accepted::<C::Error>()
.map_err(WorkerQuitReason::fatal_context( .map_err(WorkerQuitReason::fatal_context(
"process initialized notification response", "process initialized notification response",
))?; ))?;
@ -373,14 +380,14 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
)); ));
tracing::debug!("got common stream"); tracing::debug!("got common stream");
} }
Err(StreamableHttpError::SeverDoesNotSupportSse) => { Err(StreamableHttpError::ServerDoesNotSupportSse) => {
tracing::debug!("server doesn't support sse, skip common stream"); tracing::debug!("server doesn't support sse, skip common stream");
} }
Err(e) => { Err(e) => {
// fail to get common stream // fail to get common stream
tracing::error!("fail to get common stream: {e}"); tracing::error!("fail to get common stream: {e}");
return Err(WorkerQuitReason::fatal( return Err(WorkerQuitReason::fatal(
"fail to get general purpose event stream", e,
"get general purpose event stream", "get general purpose event stream",
)); ));
} }

58
crates/rmcp/src/transport/streamable_http_server/session/local.rs
View file

@ -296,12 +296,8 @@ pub enum SessionError {
SessionServiceTerminated, SessionServiceTerminated,
#[error("Invalid event id")] #[error("Invalid event id")]
InvalidEventId, InvalidEventId,
#[error("Transport closed")]
TransportClosed,
#[error("IO error: {0}")] #[error("IO error: {0}")]
Io(#[from] std::io::Error), Io(#[from] std::io::Error),
#[error("Tokio join error {0}")]
TokioJoinError(#[from] tokio::task::JoinError),
} }
impl From<SessionError> for std::io::Error { impl From<SessionError> for std::io::Error {
@ -317,7 +313,7 @@ enum OutboundChannel {
RequestWise { id: HttpRequestId, close: bool }, RequestWise { id: HttpRequestId, close: bool },
Common, Common,
} }
#[derive(Debug)]
pub struct StreamableHttpMessageReceiver { pub struct StreamableHttpMessageReceiver {
pub http_request_id: Option<HttpRequestId>, pub http_request_id: Option<HttpRequestId>,
pub inner: Receiver<ServerSseMessage>, pub inner: Receiver<ServerSseMessage>,
@ -534,8 +530,8 @@ impl LocalSessionWorker {
} }
} }
} }
#[derive(Debug)]
enum SessionEvent { pub enum SessionEvent {
ClientMessage { ClientMessage {
message: ClientJsonRpcMessage, message: ClientJsonRpcMessage,
http_request_id: Option<HttpRequestId>, http_request_id: Option<HttpRequestId>,
@ -695,14 +691,31 @@ impl LocalSessionHandle {
pub type SessionTransport = WorkerTransport<LocalSessionWorker>; pub type SessionTransport = WorkerTransport<LocalSessionWorker>;
#[derive(Debug, Error)]
pub enum LocalSessionWorkerError {
#[error("transport terminated")]
TransportTerminated,
#[error("unexpected message: {0:?}")]
UnexpectedEvent(SessionEvent),
#[error("fail to send initialize request {0}")]
FailToSendInitializeRequest(SessionError),
#[error("fail to handle message: {0}")]
FailToHandleMessage(SessionError),
#[error("keep alive timeout after {}ms", _0.as_millis())]
KeepAliveTimeout(Duration),
#[error("Transport closed")]
TransportClosed,
#[error("Tokio join error {0}")]
TokioJoinError(#[from] tokio::task::JoinError),
}
impl Worker for LocalSessionWorker { impl Worker for LocalSessionWorker {
type Error = SessionError; type Error = LocalSessionWorkerError;
type Role = RoleServer; type Role = RoleServer;
fn err_closed() -> Self::Error { fn err_closed() -> Self::Error {
SessionError::TransportClosed LocalSessionWorkerError::TransportClosed
} }
fn err_join(e: tokio::task::JoinError) -> Self::Error { fn err_join(e: tokio::task::JoinError) -> Self::Error {
SessionError::TokioJoinError(e) LocalSessionWorkerError::TokioJoinError(e)
} }
fn config(&self) -> crate::transport::worker::WorkerConfig { fn config(&self) -> crate::transport::worker::WorkerConfig {
crate::transport::worker::WorkerConfig { crate::transport::worker::WorkerConfig {
@ -711,18 +724,24 @@ impl Worker for LocalSessionWorker {
} }
} }
#[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))] #[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> { async fn run(
mut self,
mut context: WorkerContext<Self>,
) -> Result<(), WorkerQuitReason<Self::Error>> {
enum InnerEvent { enum InnerEvent {
FromHttpService(SessionEvent), FromHttpService(SessionEvent),
FromHandler(WorkerSendRequest<LocalSessionWorker>), FromHandler(WorkerSendRequest<LocalSessionWorker>),
} }
// waiting for initialize request // waiting for initialize request
let evt = self.event_rx.recv().await.ok_or_else(|| { let evt = self.event_rx.recv().await.ok_or_else(|| {
WorkerQuitReason::fatal("transport terminated", "get initialize request") WorkerQuitReason::fatal(
LocalSessionWorkerError::TransportTerminated,
"get initialize request",
)
})?; })?;
let SessionEvent::InitializeRequest { request, responder } = evt else { let SessionEvent::InitializeRequest { request, responder } = evt else {
return Err(WorkerQuitReason::fatal( return Err(WorkerQuitReason::fatal(
"unexpected message", LocalSessionWorkerError::UnexpectedEvent(evt),
"get initialize request", "get initialize request",
)); ));
}; };
@ -732,7 +751,9 @@ impl Worker for LocalSessionWorker {
.send(Ok(send_initialize_response.message)) .send(Ok(send_initialize_response.message))
.map_err(|_| { .map_err(|_| {
WorkerQuitReason::fatal( WorkerQuitReason::fatal(
"failed to send initialize response to http service", LocalSessionWorkerError::FailToSendInitializeRequest(
SessionError::SessionServiceTerminated,
),
"send initialize response", "send initialize response",
) )
})?; })?;
@ -749,7 +770,7 @@ impl Worker for LocalSessionWorker {
if let Some(event) = event { if let Some(event) = event {
InnerEvent::FromHttpService(event) InnerEvent::FromHttpService(event)
} else { } else {
return Err(WorkerQuitReason::fatal("session dropped", "waiting next session event")) return Err(WorkerQuitReason::fatal(LocalSessionWorkerError::TransportTerminated, "waiting next session event"))
} }
}, },
from_handler = context.recv_from_handler() => { from_handler = context.recv_from_handler() => {
@ -759,7 +780,7 @@ impl Worker for LocalSessionWorker {
return Err(WorkerQuitReason::Cancelled) return Err(WorkerQuitReason::Cancelled)
} }
_ = keep_alive_timeout => { _ = keep_alive_timeout => {
return Err(WorkerQuitReason::fatal("keep live timeout", "poll next session event")) return Err(WorkerQuitReason::fatal(LocalSessionWorkerError::KeepAliveTimeout(keep_alive), "poll next session event"))
} }
}; };
match event { match event {
@ -779,7 +800,10 @@ impl Worker for LocalSessionWorker {
// no need to unregister resource // no need to unregister resource
} }
}; };
let handle_result = self.handle_server_message(message).await; let handle_result = self
.handle_server_message(message)
.await
.map_err(LocalSessionWorkerError::FailToHandleMessage);
let _ = responder.send(handle_result).inspect_err(|error| { let _ = responder.send(handle_result).inspect_err(|error| {
tracing::warn!(?error, "failed to send message to http service handler"); tracing::warn!(?error, "failed to send message to http service handler");
}); });

26
crates/rmcp/src/transport/worker.rs
View file

@ -7,12 +7,12 @@ use super::{IntoTransport, Transport};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage}; use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
#[derive(Debug, thiserror::Error)] #[derive(Debug, thiserror::Error)]
pub enum WorkerQuitReason { pub enum WorkerQuitReason<E> {
#[error("Join error {0}")] #[error("Join error {0}")]
Join(#[from] tokio::task::JoinError), Join(#[from] tokio::task::JoinError),
#[error("Transport fatal {error}, when {context}")] #[error("Transport fatal {error}, when {context}")]
Fatal { Fatal {
error: Cow<'static, str>, error: E,
context: Cow<'static, str>, context: Cow<'static, str>,
}, },
#[error("Transport canncelled")] #[error("Transport canncelled")]
@ -23,18 +23,16 @@ pub enum WorkerQuitReason {
HandlerTerminated, HandlerTerminated,
} }
impl WorkerQuitReason { impl<E: std::error::Error + Send + 'static> WorkerQuitReason<E> {
pub fn fatal(msg: impl Into<Cow<'static, str>>, context: impl Into<Cow<'static, str>>) -> Self { pub fn fatal(error: E, context: impl Into<Cow<'static, str>>) -> Self {
Self::Fatal { Self::Fatal {
error: msg.into(), error,
context: context.into(), context: context.into(),
} }
} }
pub fn fatal_context<E: std::error::Error>( pub fn fatal_context(context: impl Into<Cow<'static, str>>) -> impl FnOnce(E) -> Self {
context: impl Into<Cow<'static, str>>,
) -> impl FnOnce(E) -> Self {
|e| Self::Fatal { |e| Self::Fatal {
error: Cow::Owned(format!("{e}")), error: e,
context: context.into(), context: context.into(),
} }
} }
@ -48,7 +46,7 @@ pub trait Worker: Sized + Send + 'static {
fn run( fn run(
self, self,
context: WorkerContext<Self>, context: WorkerContext<Self>,
) -> impl Future<Output = Result<(), WorkerQuitReason>> + Send; ) -> impl Future<Output = Result<(), WorkerQuitReason<Self::Error>>> + Send;
fn config(&self) -> WorkerConfig { fn config(&self) -> WorkerConfig {
WorkerConfig::default() WorkerConfig::default()
} }
@ -62,7 +60,7 @@ pub struct WorkerSendRequest<W: Worker> {
pub struct WorkerTransport<W: Worker> { pub struct WorkerTransport<W: Worker> {
rx: tokio::sync::mpsc::Receiver<RxJsonRpcMessage<W::Role>>, rx: tokio::sync::mpsc::Receiver<RxJsonRpcMessage<W::Role>>,
send_service: tokio::sync::mpsc::Sender<WorkerSendRequest<W>>, send_service: tokio::sync::mpsc::Sender<WorkerSendRequest<W>>,
join_handle: Option<tokio::task::JoinHandle<Result<(), WorkerQuitReason>>>, join_handle: Option<tokio::task::JoinHandle<Result<(), WorkerQuitReason<W::Error>>>>,
_drop_guard: tokio_util::sync::DropGuard, _drop_guard: tokio_util::sync::DropGuard,
ct: CancellationToken, ct: CancellationToken,
} }
@ -159,14 +157,16 @@ impl<W: Worker> WorkerContext<W> {
pub async fn send_to_handler( pub async fn send_to_handler(
&mut self, &mut self,
item: RxJsonRpcMessage<W::Role>, item: RxJsonRpcMessage<W::Role>,
) -> Result<(), WorkerQuitReason> { ) -> Result<(), WorkerQuitReason<W::Error>> {
self.to_handler_tx self.to_handler_tx
.send(item) .send(item)
.await .await
.map_err(|_| WorkerQuitReason::HandlerTerminated) .map_err(|_| WorkerQuitReason::HandlerTerminated)
} }
pub async fn recv_from_handler(&mut self) -> Result<WorkerSendRequest<W>, WorkerQuitReason> { pub async fn recv_from_handler(
&mut self,
) -> Result<WorkerSendRequest<W>, WorkerQuitReason<W::Error>> {
self.from_handler_rx self.from_handler_rx
.recv() .recv()
.await .await