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

fix: propagate tracing spans when spawning new tokio tasks (#334)

Browse source
This commit is contained in:
Nate Usher 2025-07-28 11:39:42 -07:00 • committed by GitHub
commit 179296134a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 16 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -30,7 +30,7 @@ use tokio_util::sync::{CancellationToken, DropGuard};
#[cfg(feature = "tower")] #[cfg(feature = "tower")]
#[cfg_attr(docsrs, doc(cfg(feature = "tower")))] #[cfg_attr(docsrs, doc(cfg(feature = "tower")))]
pub use tower::*; pub use tower::*;
use tracing::instrument; use tracing::{Instrument as _, instrument};
#[derive(Error, Debug)] #[derive(Error, Debug)]
#[non_exhaustive] #[non_exhaustive]
pub enum ServiceError { pub enum ServiceError {
@ -573,6 +573,7 @@ where
// let mut stream = std::pin::pin!(stream); // let mut stream = std::pin::pin!(stream);
let serve_loop_ct = ct.child_token(); let serve_loop_ct = ct.child_token();
let peer_return: Peer<R> = peer.clone(); let peer_return: Peer<R> = peer.clone();
let current_span = tracing::Span::current();
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
let mut transport = transport.into_transport(); let mut transport = transport.into_transport();
let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new(); let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new();
@ -687,12 +688,13 @@ where
ct.cancel(); ct.cancel();
} }
let send = transport.send(m); let send = transport.send(m);
let current_span = tracing::Span::current();
tokio::spawn(async move { tokio::spawn(async move {
let send_result = send.await; let send_result = send.await;
if let Err(error) = send_result { if let Err(error) = send_result {
tracing::error!(%error, "fail to response message"); tracing::error!(%error, "fail to response message");
} }
}); }.instrument(current_span));
} }
} }
Event::ProxyMessage(PeerSinkMessage::Request { Event::ProxyMessage(PeerSinkMessage::Request {
@ -704,10 +706,11 @@ where
let send = transport.send(JsonRpcMessage::request(request, id.clone())); let send = transport.send(JsonRpcMessage::request(request, id.clone()));
{ {
let id = id.clone(); let id = id.clone();
let current_span = tracing::Span::current();
send_task_set.spawn(send.map(move |r| SendTaskResult::Request { send_task_set.spawn(send.map(move |r| SendTaskResult::Request {
id, id,
result: r.map_err(DynamicTransportError::new::<T, R>), result: r.map_err(DynamicTransportError::new::<T, R>),
})); }).instrument(current_span));
} }
} }
Event::ProxyMessage(PeerSinkMessage::Notification { Event::ProxyMessage(PeerSinkMessage::Notification {
@ -724,11 +727,12 @@ where
Err(notification) => notification, Err(notification) => notification,
}; };
let send = transport.send(JsonRpcMessage::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 { send_task_set.spawn(send.map(move |result| SendTaskResult::Notification {
responder, responder,
cancellation_param, cancellation_param,
result: result.map_err(DynamicTransportError::new::<T, R>), result: result.map_err(DynamicTransportError::new::<T, R>),
})); }).instrument(current_span));
} }
Event::PeerMessage(JsonRpcMessage::Request(JsonRpcRequest { Event::PeerMessage(JsonRpcMessage::Request(JsonRpcRequest {
id, id,
@ -755,8 +759,11 @@ where
meta, meta,
extensions, extensions,
}; };
let current_span = tracing::Span::current();
tokio::spawn(async move { tokio::spawn(async move {
let result = service.handle_request(request, context).await; let result = service
.handle_request(request, context)
.await;
let response = match result { let response = match result {
Ok(result) => { Ok(result) => {
tracing::debug!(%id, ?result, "response message"); tracing::debug!(%id, ?result, "response message");
@ -768,7 +775,7 @@ where
} }
}; };
let _send_result = sink.send(response).await; let _send_result = sink.send(response).await;
}); }.instrument(current_span));
} }
} }
Event::PeerMessage(JsonRpcMessage::Notification(JsonRpcNotification { Event::PeerMessage(JsonRpcMessage::Notification(JsonRpcNotification {
@ -799,12 +806,13 @@ where
meta, meta,
extensions, extensions,
}; };
let current_span = tracing::Span::current();
tokio::spawn(async move { tokio::spawn(async move {
let result = service.handle_notification(notification, context).await; let result = service.handle_notification(notification, context).await;
if let Err(error) = result { if let Err(error) = result {
tracing::warn!(%error, "Error sending notification"); tracing::warn!(%error, "Error sending notification");
} }
}); }.instrument(current_span));
} }
} }
Event::PeerMessage(JsonRpcMessage::Response(JsonRpcResponse { Event::PeerMessage(JsonRpcMessage::Response(JsonRpcResponse {
@ -849,7 +857,7 @@ where
} }
tracing::info!(?quit_reason, "serve finished"); tracing::info!(?quit_reason, "serve finished");
quit_reason quit_reason
}); }.instrument(current_span));
RunningService { RunningService {
service, service,
peer: peer_return, peer: peer_return,