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

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
1 changed files with 71 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -296,12 +296,8 @@ pub enum SessionError {
SessionServiceTerminated,
#[error("Invalid event id")]
InvalidEventId,
#[error("Transport closed")]
TransportClosed,
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Tokio join error {0}")]
TokioJoinError(#[from] tokio::task::JoinError),
}
impl From<SessionError> for std::io::Error {
@ -317,7 +313,7 @@ enum OutboundChannel {
RequestWise { id: HttpRequestId, close: bool },
Common,
}
#[derive(Debug)]
pub struct StreamableHttpMessageReceiver {
pub http_request_id: Option<HttpRequestId>,
pub inner: Receiver<ServerSseMessage>,
@ -534,8 +530,8 @@ impl LocalSessionWorker {
}
}
}
enum SessionEvent {
#[derive(Debug)]
pub enum SessionEvent {
ClientMessage {
message: ClientJsonRpcMessage,
http_request_id: Option<HttpRequestId>,
@ -695,14 +691,31 @@ impl LocalSessionHandle {
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 {
type Error = SessionError;
type Error = LocalSessionWorkerError;
type Role = RoleServer;
fn err_closed() -> Self::Error {
SessionError::TransportClosed
LocalSessionWorkerError::TransportClosed
}
fn err_join(e: tokio::task::JoinError) -> Self::Error {
SessionError::TokioJoinError(e)
LocalSessionWorkerError::TokioJoinError(e)
}
fn config(&self) -> 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()))]
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 {
FromHttpService(SessionEvent),
FromHandler(WorkerSendRequest<LocalSessionWorker>),
}
// waiting for initialize request
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 {
return Err(WorkerQuitReason::fatal(
"unexpected message",
LocalSessionWorkerError::UnexpectedEvent(evt),
"get initialize request",
));
};
@ -732,7 +751,9 @@ impl Worker for LocalSessionWorker {
.send(Ok(send_initialize_response.message))
.map_err(|_| {
WorkerQuitReason::fatal(
"failed to send initialize response to http service",
LocalSessionWorkerError::FailToSendInitializeRequest(
SessionError::SessionServiceTerminated,
),
"send initialize response",
)
})?;
@ -749,7 +770,7 @@ impl Worker for LocalSessionWorker {
if let Some(event) = event {
InnerEvent::FromHttpService(event)
} 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() => {
@ -759,7 +780,7 @@ impl Worker for LocalSessionWorker {
return Err(WorkerQuitReason::Cancelled)
}
_ = 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 {
@ -779,7 +800,10 @@ impl Worker for LocalSessionWorker {
// 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| {
tracing::warn!(?error, "failed to send message to http service handler");
});