refactor: provide http server as tower service (#228)

* refactor: streamable http server as tower service
This commit is contained in:
4t145 2025-05-28 19:21:40 +08:00 • committed by GitHub
commit 209be7bdb8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 1828 additions and 1180 deletions

View file

@ -39,7 +39,9 @@ reqwest = { version = "0.12", default-features = false, features = [
"json",
"stream",
], optional = true }
sse-stream = { version = "0.2.0", optional = true }
sse-stream = { version = "0.2", optional = true }
http = { version = "1", optional = true }
url = { version = "2.4", optional = true }
@ -57,7 +59,9 @@ axum = { version = "0.8", features = [], optional = true }
rand = { version = "0.9", optional = true }
tokio-stream = { version = "0.1", optional = true }
uuid = { version = "1", features = ["v4"], optional = true }
http-body = { version = "1", optional = true }
http-body-util = { version = "0.1", optional = true }
bytes = { version = "1", optional = true }
# macro
rmcp-macros = { version = "0.1", workspace = true, optional = true }
@ -74,7 +78,17 @@ reqwest = ["__reqwest", "reqwest?/rustls-tls"]
reqwest-tls-no-provider = ["__reqwest", "reqwest?/rustls-tls-no-provider"]
axum = ["dep:axum"]
server-side-http = [
"uuid",
"dep:rand",
"dep:tokio-stream",
"dep:http",
"dep:http-body",
"dep:http-body-util",
"dep:bytes",
"dep:sse-stream",
"tower",
]
# SSE client
client-side-sse = ["dep:sse-stream", "dep:http"]
@ -97,15 +111,12 @@ transport-child-process = [
transport-sse-server = [
"transport-async-rw",
"transport-worker",
"axum",
"dep:rand",
"dep:tokio-stream",
"uuid",
"server-side-http",
"dep:axum",
]
transport-streamable-http-server = [
"transport-streamable-http-server-session",
"axum",
"uuid",
"server-side-http",
]
transport-streamable-http-server-session = [
"transport-async-rw",

View file

@ -488,7 +488,7 @@ pub struct RequestContext<R: ServiceRole> {
}
/// Use this function to skip initialization process
pub async fn serve_directly<R, S, T, E, A>(
pub fn serve_directly<R, S, T, E, A>(
service: S,
transport: T,
peer_info: Option<R::PeerInfo>,
@ -499,11 +499,11 @@ where
T: IntoTransport<R, E, A>,
E: std::error::Error + Send + Sync + 'static,
{
serve_directly_with_ct(service, transport, peer_info, Default::default()).await
serve_directly_with_ct(service, transport, peer_info, Default::default())
}
/// Use this function to skip initialization process
pub async fn serve_directly_with_ct<R, S, T, E, A>(
pub fn serve_directly_with_ct<R, S, T, E, A>(
service: S,
transport: T,
peer_info: Option<R::PeerInfo>,
@ -516,11 +516,11 @@ where
E: std::error::Error + Send + Sync + 'static,
{
let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
serve_inner(service, transport, peer, peer_rx, ct).await
serve_inner(service, transport, peer, peer_rx, ct)
}
#[instrument(skip_all)]
async fn serve_inner<R, S, T, E, A>(
fn serve_inner<R, S, T, E, A>(
service: S,
transport: T,
peer: Peer<R>,

View file

@ -111,7 +111,7 @@ pub async fn serve_client<S, T, E, A>(
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
E: std::error::Error + Send + Sync + 'static,
{
serve_client_with_ct(service, transport, Default::default()).await
}
@ -124,7 +124,7 @@ pub async fn serve_client_with_ct<S, T, E, A>(
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
E: std::error::Error + Send + Sync + 'static,
{
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
@ -175,7 +175,7 @@ where
context: "send initialized notification".into(),
})?;
let (peer, peer_rx) = Peer::new(id_provider, Some(initialize_result));
Ok(serve_inner(service, transport, peer, peer_rx, ct).await)
Ok(serve_inner(service, transport, peer, peer_rx, ct))
}
macro_rules! method {

View file

@ -70,7 +70,7 @@ impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
) -> 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,
E: std::error::Error + Send + Sync + 'static,
Self: Sized,
{
serve_server_with_ct(self, transport, ct)
@ -84,7 +84,7 @@ pub async fn serve_server<S, T, E, A>(
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
E: std::error::Error + Send + Sync + 'static,
{
serve_server_with_ct(service, transport, CancellationToken::new()).await
}
@ -143,7 +143,7 @@ pub async fn serve_server_with_ct<S, T, E, A>(
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
E: std::error::Error + Send + Sync + 'static,
{
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
@ -212,7 +212,7 @@ where
};
let _ = service.handle_notification(notification).await;
// Continue processing service
Ok(serve_inner(service, transport, peer, peer_rx, ct).await)
Ok(serve_inner(service, transport, peer, peer_rx, ct))
}
macro_rules! method {

View file

@ -7,7 +7,7 @@
//! | transport | client | server |
//! |:-: |:-: |:-: |
//! | std IO | [`child_process::TokioChildProcess`] | [`io::stdio`] |
//! | streamable http | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::session::create_session`] |
//! | streamable http | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::StreamableHttpService`] |
//! | sse | [`sse_client::SseClientTransport`] | [`sse_server::SseServer`] |
//!
//!## Helper Transport Types
@ -64,6 +64,8 @@
//! }
//! ```
use std::sync::Arc;
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
pub mod sink_stream;
@ -122,7 +124,7 @@ pub use auth::{AuthError, AuthorizationManager, AuthorizationSession, Authorized
pub mod streamable_http_server;
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub use streamable_http_server::axum::StreamableHttpServer;
pub use streamable_http_server::tower::{StreamableHttpServerConfig, StreamableHttpService};
#[cfg(feature = "transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client")))]
@ -138,7 +140,7 @@ pub trait Transport<R>: Send
where
R: ServiceRole,
{
type Error;
type Error: std::error::Error + Send + Sync + 'static;
/// Send a message to the transport
///
/// Notice that the future returned by this function should be `Send` and `'static`.
@ -169,9 +171,73 @@ impl<R, T, E> IntoTransport<R, E, TransportAdapterIdentity> for T
where
T: Transport<R, Error = E> + Send + 'static,
R: ServiceRole,
E: std::error::Error + Send + 'static,
E: std::error::Error + Send + Sync + 'static,
{
fn into_transport(self) -> impl Transport<R, Error = E> + 'static {
self
}
}
/// A transport that can send a single message and then close itself
pub struct OneshotTransport<R>
where
R: ServiceRole,
{
message: Option<RxJsonRpcMessage<R>>,
sender: tokio::sync::mpsc::Sender<TxJsonRpcMessage<R>>,
finished_signal: Arc<tokio::sync::Notify>,
}
impl<R> OneshotTransport<R>
where
R: ServiceRole,
{
pub fn new(
message: RxJsonRpcMessage<R>,
) -> (Self, tokio::sync::mpsc::Receiver<TxJsonRpcMessage<R>>) {
let (sender, receiver) = tokio::sync::mpsc::channel(16);
(
Self {
message: Some(message),
sender,
finished_signal: Arc::new(tokio::sync::Notify::new()),
},
receiver,
)
}
}
impl<R> Transport<R> for OneshotTransport<R>
where
R: ServiceRole,
{
type Error = tokio::sync::mpsc::error::SendError<TxJsonRpcMessage<R>>;
fn send(
&mut self,
item: TxJsonRpcMessage<R>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
let sender = self.sender.clone();
let terminate = matches!(item, TxJsonRpcMessage::<R>::Response(_));
let signal = self.finished_signal.clone();
async move {
sender.send(item).await?;
if terminate {
signal.notify_waiters();
}
Ok(())
}
}
async fn receive(&mut self) -> Option<RxJsonRpcMessage<R>> {
if self.message.is_none() {
self.finished_signal.notified().await;
}
self.message.take()
}
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {
self.message.take();
std::future::ready(Ok(()))
}
}

View file

@ -2,7 +2,7 @@
feature = "transport-streamable-http-server",
feature = "transport-sse-server"
))]
pub mod axum;
pub mod server_side_http;
pub mod http_header;

View file

@ -1,9 +0,0 @@
use std::{sync::Arc, time::Duration};
pub type SessionId = Arc<str>;
pub fn session_id() -> SessionId {
uuid::Uuid::new_v4().to_string().into()
}
pub const DEFAULT_AUTO_PING_INTERVAL: Duration = Duration::from_secs(15);

View file

@ -0,0 +1,140 @@
use std::{convert::Infallible, fmt::Display, sync::Arc, time::Duration};
use bytes::{Buf, Bytes};
use http::Response;
use http_body::Body;
use http_body_util::{BodyExt, Empty, Full, combinators::UnsyncBoxBody};
use sse_stream::{KeepAlive, Sse, SseBody};
use super::http_header::EVENT_STREAM_MIME_TYPE;
use crate::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
pub type SessionId = Arc<str>;
pub fn session_id() -> SessionId {
uuid::Uuid::new_v4().to_string().into()
}
pub const DEFAULT_AUTO_PING_INTERVAL: Duration = Duration::from_secs(15);
pub(crate) type BoxResponse = Response<UnsyncBoxBody<Bytes, Infallible>>;
pub(crate) fn accepted_response() -> Response<UnsyncBoxBody<Bytes, Infallible>> {
Response::builder()
.status(http::StatusCode::ACCEPTED)
.body(Empty::new().boxed_unsync())
.expect("valid response")
}
pin_project_lite::pin_project! {
struct TokioTimer {
#[pin]
sleep: tokio::time::Sleep,
}
}
impl Future for TokioTimer {
type Output = ();
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let this = self.project();
this.sleep.poll(cx)
}
}
impl sse_stream::Timer for TokioTimer {
fn from_duration(duration: Duration) -> Self {
Self {
sleep: tokio::time::sleep(duration),
}
}
fn reset(self: std::pin::Pin<&mut Self>, when: std::time::Instant) {
let this = self.project();
this.sleep.reset(tokio::time::Instant::from_std(when));
}
}
#[derive(Debug, Clone)]
pub struct ServerSseMessage {
pub event_id: Option<String>,
pub message: Arc<ServerJsonRpcMessage>,
}
pub(crate) fn sse_stream_response(
stream: impl futures::Stream<Item = ServerSseMessage> + Send + 'static,
keep_alive: Option<Duration>,
) -> Response<UnsyncBoxBody<Bytes, Infallible>> {
use futures::StreamExt;
let stream = SseBody::new(stream.map(|message| {
let data = serde_json::to_string(&message.message).expect("valid message");
let mut sse = Sse::default().data(data);
sse.id = message.event_id;
Result::<Sse, Infallible>::Ok(sse)
}));
let stream = match keep_alive {
Some(duration) => stream
.with_keep_alive::<TokioTimer>(KeepAlive::new().interval(duration))
.boxed_unsync(),
None => stream.boxed_unsync(),
};
Response::builder()
.status(http::StatusCode::OK)
.header(http::header::CONTENT_TYPE, EVENT_STREAM_MIME_TYPE)
.header(http::header::CACHE_CONTROL, "no-cache")
.body(stream)
.expect("valid response")
}
pub(crate) const fn internal_error_response<E: Display>(
context: &str,
) -> impl FnOnce(E) -> Response<UnsyncBoxBody<Bytes, Infallible>> {
move |error| {
tracing::error!("Internal server error when {context}: {error}");
Response::builder()
.status(http::StatusCode::INTERNAL_SERVER_ERROR)
.body(
Full::new(Bytes::from(format!(
"Encounter an error when {context}: {error}"
)))
.boxed_unsync(),
)
.expect("valid response")
}
}
pub(crate) async fn expect_json<B>(
body: B,
) -> Result<ClientJsonRpcMessage, Response<UnsyncBoxBody<Bytes, Infallible>>>
where
B: Body + Send + 'static,
B::Error: Display,
{
match body.collect().await {
Ok(bytes) => {
match serde_json::from_reader::<_, ClientJsonRpcMessage>(bytes.aggregate().reader()) {
Ok(message) => Ok(message),
Err(e) => {
let response = Response::builder()
.status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
.body(
Full::new(Bytes::from(format!("fail to deserialize request body {e}")))
.boxed_unsync(),
)
.expect("valid response");
Err(response)
}
}
}
Err(e) => {
let response = Response::builder()
.status(http::StatusCode::INTERNAL_SERVER_ERROR)
.body(
Full::new(Bytes::from(format!("Failed to read request body: {e}")))
.boxed_unsync(),
)
.expect("valid response");
Err(response)
}
}
}

View file

@ -24,6 +24,7 @@ impl<Role: ServiceRole, Si, St> Transport<Role> for SinkStreamTransport<Si, St>
where
St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin,
Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
Si::Error: std::error::Error + Send + Sync + 'static,
{
type Error = Si::Error;
@ -56,7 +57,7 @@ where
Role: ServiceRole,
Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin + 'static,
Si::Error: std::error::Error + Send + 'static,
Si::Error: std::error::Error + Send + Sync + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = Si::Error> + 'static {
SinkStreamTransport::new(self.0, self.1)
@ -68,7 +69,7 @@ impl<Role, S> IntoTransport<Role, S::Error, TransportAdapterAsyncCombinedRW> for
where
Role: ServiceRole,
S: Sink<TxJsonRpcMessage<Role>> + Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
S::Error: std::error::Error + Send + 'static,
S::Error: std::error::Error + Send + Sync + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = S::Error> + 'static {
use futures::StreamExt;

View file

@ -19,7 +19,7 @@ use crate::{
RoleServer, Service,
model::ClientJsonRpcMessage,
service::{RxJsonRpcMessage, TxJsonRpcMessage, serve_directly_with_ct},
transport::common::axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
transport::common::server_side_http::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
};
type TxStore =
@ -313,7 +313,7 @@ impl SseServer {
let service = service_provider();
let ct = self.config.ct.child_token();
tokio::spawn(async move {
let server = serve_directly_with_ct(service, transport, None, ct).await;
let server = serve_directly_with_ct(service, transport, None, ct);
server.waiting().await?;
tokio::io::Result::Ok(())
});

View file

@ -1,5 +1,8 @@
pub mod session;
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub mod axum;
pub mod session;
pub use session::{SessionConfig, create_session};
pub mod tower;
pub use session::{SessionId, SessionManager};
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub use tower::{StreamableHttpServerConfig, StreamableHttpService};

View file

@ -1,349 +0,0 @@
use std::{
collections::HashMap,
io,
net::{Ipv6Addr, SocketAddr, SocketAddrV6},
sync::Arc,
time::Duration,
};
use axum::{
Json, Router,
extract::State,
http::{HeaderMap, HeaderValue, StatusCode, request::Parts},
response::{
IntoResponse, Response,
sse::{Event, KeepAlive, Sse},
},
routing::get,
};
use futures::Stream;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use super::session::{EventId, SessionHandle, SessionWorker, StreamableHttpMessageReceiver};
use crate::{
RoleServer, Service,
model::ClientJsonRpcMessage,
transport::common::{
axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
http_header::{HEADER_LAST_EVENT_ID, HEADER_SESSION_ID},
},
};
type SessionManager = Arc<tokio::sync::RwLock<HashMap<SessionId, SessionHandle>>>;
#[derive(Clone)]
struct App {
session_manager: SessionManager,
transport_tx: tokio::sync::mpsc::UnboundedSender<SessionWorker>,
sse_ping_interval: Duration,
}
impl App {
pub fn new(
sse_ping_interval: Duration,
) -> (Self, tokio::sync::mpsc::UnboundedReceiver<SessionWorker>) {
let (transport_tx, transport_rx) = tokio::sync::mpsc::unbounded_channel();
(
Self {
session_manager: Default::default(),
transport_tx,
sse_ping_interval,
},
transport_rx,
)
}
}
fn receiver_as_stream(
receiver: StreamableHttpMessageReceiver,
) -> impl Stream<Item = Result<Event, io::Error>> {
use futures::StreamExt;
ReceiverStream::new(receiver.inner).map(|message| {
match serde_json::to_string(&message.message) {
Ok(bytes) => Ok(Event::default()
.event("message")
.data(&bytes)
.id(message.event_id.to_string())),
Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
}
})
}
async fn post_handler(
State(app): State<App>,
parts: Parts,
Json(mut message): Json<ClientJsonRpcMessage>,
) -> Result<Response, Response> {
use futures::StreamExt;
if let Some(session_id) = parts.headers.get(HEADER_SESSION_ID).cloned() {
let session_id = session_id
.to_str()
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
tracing::debug!(session_id, ?message, "new client message");
let handle = {
let sm = app.session_manager.read().await;
let session = sm
.get(session_id)
.ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
session.clone()
};
// inject request part
message.insert_extension(parts);
match &message {
ClientJsonRpcMessage::Request(_) | ClientJsonRpcMessage::BatchRequest(_) => {
let receiver = handle.establish_request_wise_channel().await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("fail to to establish request channel: {e}"),
)
.into_response()
})?;
let http_request_id = receiver.http_request_id;
if let Err(push_err) = handle.push_message(message, http_request_id).await {
tracing::error!(session_id, ?push_err, "push message error");
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
format!("fail to push message: {push_err}"),
)
.into_response());
}
let stream =
ReceiverStream::new(receiver.inner).map(|message| match serde_json::to_string(
&message.message,
) {
Ok(bytes) => Ok(Event::default()
.event("message")
.data(&bytes)
.id(message.event_id.to_string())),
Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
});
Ok(Sse::new(stream)
.keep_alive(KeepAlive::new().interval(app.sse_ping_interval))
.into_response())
}
_ => {
let result = handle.push_message(message, None).await;
if result.is_err() {
Err((StatusCode::GONE, "session terminated").into_response())
} else {
Ok(StatusCode::ACCEPTED.into_response())
}
}
}
} else {
// expect initialize message
let session_id = session_id();
// inject request part
message.insert_extension(parts);
let (session, transport) =
super::session::create_session(session_id.clone(), Default::default());
let Ok(_) = app.transport_tx.send(transport) else {
return Err((StatusCode::GONE, "session terminated").into_response());
};
let response = session.initialize(message).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("fail to initialize: {e}"),
)
.into_response()
})?;
let mut response = Json(response).into_response();
response.headers_mut().insert(
HEADER_SESSION_ID,
HeaderValue::from_bytes(session_id.as_bytes()).expect("should be valid header value"),
);
app.session_manager
.write()
.await
.insert(session_id, session);
Ok(response)
}
}
async fn get_handler(
State(app): State<App>,
header_map: HeaderMap,
) -> Result<Sse<impl Stream<Item = Result<Event, io::Error>>>, Response> {
let session_id = header_map
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok());
if let Some(session_id) = session_id {
let last_event_id = header_map
.get(HEADER_LAST_EVENT_ID)
.and_then(|v| v.to_str().ok());
let session = {
let sm = app.session_manager.read().await;
sm.get(session_id)
.ok_or_else(|| {
(
StatusCode::NOT_FOUND,
format!("session {session_id} not found"),
)
.into_response()
})?
.clone()
};
match last_event_id {
Some(last_event_id) => {
let last_event_id = last_event_id.parse::<EventId>().map_err(|e| {
(StatusCode::BAD_REQUEST, format!("invalid event_id {e}")).into_response()
})?;
let receiver = session.resume(last_event_id).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("resume error {e}"),
)
.into_response()
})?;
let stream = receiver_as_stream(receiver);
Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(app.sse_ping_interval)))
}
None => {
let receiver = session.establish_common_channel().await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("establish common channel error {e}"),
)
.into_response()
})?;
let stream = receiver_as_stream(receiver);
Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(app.sse_ping_interval)))
}
}
} else {
Err((StatusCode::BAD_REQUEST, "missing session id").into_response())
}
}
async fn delete_handler(
State(app): State<App>,
header_map: HeaderMap,
) -> Result<StatusCode, Response> {
if let Some(session_id) = header_map.get(HEADER_SESSION_ID) {
let session_id = session_id
.to_str()
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
let session = {
let mut sm = app.session_manager.write().await;
sm.remove(session_id)
.ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?
};
session.close().await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("fail to cancel session {session_id}: tokio join error: {e}"),
)
.into_response()
})?;
tracing::debug!(session_id, "session deleted");
Ok(StatusCode::ACCEPTED)
} else {
Err((StatusCode::BAD_REQUEST, "missing session id").into_response())
}
}
#[derive(Debug, Clone)]
pub struct StreamableHttpServerConfig {
pub bind: SocketAddr,
pub path: String,
pub ct: CancellationToken,
pub sse_keep_alive: Option<Duration>,
}
impl Default for StreamableHttpServerConfig {
fn default() -> Self {
Self {
bind: SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 80, 0, 0)),
path: "/".to_string(),
ct: CancellationToken::new(),
sse_keep_alive: None,
}
}
}
#[derive(Debug)]
pub struct StreamableHttpServer {
transport_rx: tokio::sync::mpsc::UnboundedReceiver<SessionWorker>,
pub config: StreamableHttpServerConfig,
}
impl StreamableHttpServer {
pub async fn serve(bind: SocketAddr) -> io::Result<Self> {
Self::serve_with_config(StreamableHttpServerConfig {
bind,
..Default::default()
})
.await
}
pub async fn serve_with_config(config: StreamableHttpServerConfig) -> io::Result<Self> {
let (streamable_http_server, service) = Self::new(config);
let listener = tokio::net::TcpListener::bind(streamable_http_server.config.bind).await?;
let ct = streamable_http_server.config.ct.child_token();
let server = axum::serve(listener, service).with_graceful_shutdown(async move {
ct.cancelled().await;
tracing::info!("streamable http server cancelled");
});
tokio::spawn(
async move {
if let Err(e) = server.await {
tracing::error!(error = %e, "streamable http server shutdown with error");
}
}
.instrument(tracing::info_span!("streamable-http-server", bind_address = %streamable_http_server.config.bind)),
);
Ok(streamable_http_server)
}
/// Warning: This function creates a new StreamableHttpServer instance with the provided configuration.
/// `App.post_path` may be incorrect if using `Router` as an embedded router.
pub fn new(config: StreamableHttpServerConfig) -> (StreamableHttpServer, Router) {
let (app, transport_rx) =
App::new(config.sse_keep_alive.unwrap_or(DEFAULT_AUTO_PING_INTERVAL));
let router = Router::new()
.route(
&config.path,
get(get_handler).post(post_handler).delete(delete_handler),
)
.with_state(app);
let server = StreamableHttpServer {
transport_rx,
config,
};
(server, router)
}
pub fn with_service<S, F>(mut self, service_provider: F) -> CancellationToken
where
S: Service<RoleServer>,
F: Fn() -> S + Send + 'static,
{
use crate::service::ServiceExt;
let ct = self.config.ct.clone();
tokio::spawn(async move {
while let Some(transport) = self.next_transport().await {
let service = service_provider();
let ct = self.config.ct.child_token();
tokio::spawn(async move {
let server = service
.serve_with_ct(transport, ct)
.await
.map_err(tokio::io::Error::other)?;
server.waiting().await?;
tokio::io::Result::Ok(())
});
}
});
ct
}
pub fn cancel(&self) {
self.config.ct.cancel();
}
pub async fn next_transport(&mut self) -> Option<SessionWorker> {
self.transport_rx.recv().await
}
}

View file

@ -1,783 +1,54 @@
use std::{
borrow::Cow,
collections::{HashMap, HashSet, VecDeque},
num::ParseIntError,
sync::Arc,
};
use thiserror::Error;
use tokio::sync::{
mpsc::{Receiver, Sender},
oneshot,
};
use tracing::instrument;
use futures::Stream;
pub use crate::transport::common::server_side_http::SessionId;
use crate::{
RoleServer,
model::{
CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest,
JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
},
transport::{
WorkerTransport,
worker::{Worker, WorkerContext, WorkerQuitReason, WorkerSendRequest},
},
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::common::server_side_http::ServerSseMessage,
};
#[derive(Debug, Clone)]
pub struct ServerSessionMessage {
pub event_id: EventId,
pub message: Arc<ServerJsonRpcMessage>,
}
pub mod local;
pub mod never;
/// `<index>/request_id>`
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct EventId {
http_request_id: Option<HttpRequestId>,
index: usize,
}
impl std::fmt::Display for EventId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.index)?;
match &self.http_request_id {
Some(http_request_id) => write!(f, "/{http_request_id}"),
None => write!(f, ""),
}
}
}
#[derive(Debug, Clone, Error)]
pub enum EventIdParseError {
#[error("Invalid index: {0}")]
InvalidIndex(ParseIntError),
#[error("Invalid numeric request id: {0}")]
InvalidNumericRequestId(ParseIntError),
#[error("Missing request id type")]
InvalidRequestIdType,
#[error("Missing request id")]
MissingRequestId,
}
impl std::str::FromStr for EventId {
type Err = EventIdParseError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if let Some((index, request_id)) = s.split_once("/") {
let index = usize::from_str(index).map_err(EventIdParseError::InvalidIndex)?;
let request_id = u64::from_str(request_id).map_err(EventIdParseError::InvalidIndex)?;
Ok(EventId {
http_request_id: Some(request_id),
index,
})
} else {
let index = usize::from_str(s).map_err(EventIdParseError::InvalidIndex)?;
Ok(EventId {
http_request_id: None,
index,
})
}
}
}
pub use crate::transport::common::axum::SessionId;
struct CachedTx {
tx: Sender<ServerSessionMessage>,
cache: VecDeque<ServerSessionMessage>,
http_request_id: Option<HttpRequestId>,
capacity: usize,
}
impl CachedTx {
fn new(tx: Sender<ServerSessionMessage>, http_request_id: Option<HttpRequestId>) -> Self {
Self {
cache: VecDeque::with_capacity(tx.capacity()),
capacity: tx.capacity(),
tx,
http_request_id,
}
}
fn new_common(tx: Sender<ServerSessionMessage>) -> Self {
Self::new(tx, None)
}
async fn send(&mut self, message: ServerJsonRpcMessage) {
let index = self.cache.back().map_or(0, |m| m.event_id.index + 1);
let event_id = EventId {
http_request_id: self.http_request_id,
index,
};
let message = ServerSessionMessage {
event_id: event_id.clone(),
message: Arc::new(message),
};
if self.cache.len() >= self.capacity {
self.cache.pop_front();
self.cache.push_back(message.clone());
} else {
self.cache.push_back(message.clone());
}
let _ = self.tx.send(message).await.inspect_err(|e| {
let event_id = &e.0.event_id;
tracing::trace!(%event_id, "trying to send message in a closed session")
});
}
async fn sync(&mut self, index: usize) -> Result<(), SessionError> {
let Some(front) = self.cache.front() else {
return Ok(());
};
let sync_index = index.saturating_sub(front.event_id.index);
if sync_index > self.cache.len() {
// invalid index
return Err(SessionError::InvalidEventId);
}
for message in self.cache.iter().skip(sync_index) {
let send_result = self.tx.send(message.clone()).await;
if send_result.is_err() {
return Err(SessionError::ChannelClosed(
message.event_id.http_request_id,
));
}
}
Ok(())
}
}
struct HttpRequestWise {
resources: HashSet<ResourceKey>,
tx: CachedTx,
}
type HttpRequestId = u64;
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
enum ResourceKey {
McpRequestId(RequestId),
ProgressToken(ProgressToken),
}
pub struct SessionWorker {
id: SessionId,
next_http_request_id: HttpRequestId,
tx_router: HashMap<HttpRequestId, HttpRequestWise>,
resource_router: HashMap<ResourceKey, HttpRequestId>,
common: CachedTx,
event_rx: Receiver<SessionEvent>,
session_config: SessionConfig,
}
impl SessionWorker {
pub fn id(&self) -> &SessionId {
&self.id
}
}
#[derive(Debug, Error)]
pub enum SessionError {
#[error("Invalid request id: {0}")]
DuplicatedRequestId(HttpRequestId),
#[error("Channel closed: {0:?}")]
ChannelClosed(Option<HttpRequestId>),
#[error("Cannot parse event id: {0}")]
EventIdParseError(Cow<'static, str>),
#[error("Session service terminated")]
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 {
fn from(value: SessionError) -> Self {
match value {
SessionError::Io(io) => io,
_ => std::io::Error::new(std::io::ErrorKind::Other, format!("Session error: {value}")),
}
}
}
enum OutboundChannel {
RequestWise { id: HttpRequestId, close: bool },
Common,
}
pub struct StreamableHttpMessageReceiver {
pub http_request_id: Option<HttpRequestId>,
pub inner: Receiver<ServerSessionMessage>,
}
impl SessionWorker {
fn unregister_resource(&mut self, resource: &ResourceKey) {
if let Some(http_request_id) = self.resource_router.remove(resource) {
tracing::trace!(?resource, http_request_id, "unregister resource");
if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
channel.resources.remove(resource);
if channel.resources.is_empty() {
tracing::debug!(http_request_id, "close http request wise channel");
self.tx_router.remove(&http_request_id);
}
}
}
}
fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
tracing::trace!(?resource, http_request_id, "register resource");
if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
channel.resources.insert(resource.clone());
self.resource_router.insert(resource, http_request_id);
}
}
fn register_request(
&mut self,
request: &JsonRpcRequest<ClientRequest>,
http_request_id: HttpRequestId,
) {
use crate::model::GetMeta;
self.register_resource(
ResourceKey::McpRequestId(request.id.clone()),
http_request_id,
);
if let Some(progress_token) = request.request.get_meta().get_progress_token() {
self.register_resource(
ResourceKey::ProgressToken(progress_token.clone()),
http_request_id,
);
}
}
fn catch_cancellation_notification(
&mut self,
notification: &JsonRpcNotification<ClientNotification>,
) {
if let ClientNotification::CancelledNotification(n) = &notification.notification {
let request_id = n.params.request_id.clone();
let resource = ResourceKey::McpRequestId(request_id);
self.unregister_resource(&resource);
}
}
fn next_http_request_id(&mut self) -> HttpRequestId {
let id = self.next_http_request_id;
self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
id
}
async fn establish_request_wise_channel(
&mut self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let http_request_id = self.next_http_request_id();
let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
self.tx_router.insert(
http_request_id,
HttpRequestWise {
resources: Default::default(),
tx: CachedTx::new(tx, Some(http_request_id)),
},
);
tracing::debug!(http_request_id, "establish new request wise channel");
Ok(StreamableHttpMessageReceiver {
http_request_id: Some(http_request_id),
inner: rx,
})
}
fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel {
match &message {
ServerJsonRpcMessage::Request(_) => OutboundChannel::Common,
ServerJsonRpcMessage::Notification(JsonRpcNotification {
notification:
ServerNotification::ProgressNotification(Notification {
params: ProgressNotificationParam { progress_token, .. },
..
}),
..
}) => {
let id = self
.resource_router
.get(&ResourceKey::ProgressToken(progress_token.clone()));
if let Some(id) = id {
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::Notification(JsonRpcNotification {
notification:
ServerNotification::CancelledNotification(Notification {
params: CancelledNotificationParam { request_id, .. },
..
}),
..
}) => {
if let Some(id) = self
.resource_router
.get(&ResourceKey::McpRequestId(request_id.clone()))
{
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::Notification(_) => OutboundChannel::Common,
ServerJsonRpcMessage::Response(json_rpc_response) => {
if let Some(id) = self
.resource_router
.get(&ResourceKey::McpRequestId(json_rpc_response.id.clone()))
{
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::Error(json_rpc_error) => {
if let Some(id) = self
.resource_router
.get(&ResourceKey::McpRequestId(json_rpc_error.id.clone()))
{
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::BatchRequest(_) | ServerJsonRpcMessage::BatchResponse(_) => {
// the server side should never yield a batch request or response now
unreachable!("server side won't yield batch request or response")
}
}
}
async fn handle_server_message(
&mut self,
message: ServerJsonRpcMessage,
) -> Result<(), SessionError> {
let outbound_channel = self.resolve_outbound_channel(&message);
match outbound_channel {
OutboundChannel::RequestWise { id, close } => {
if let Some(request_wise) = self.tx_router.get_mut(&id) {
request_wise.tx.send(message).await;
if close {
self.tx_router.remove(&id);
}
} else {
return Err(SessionError::ChannelClosed(Some(id)));
}
}
OutboundChannel::Common => self.common.send(message).await,
}
Ok(())
}
async fn resume(
&mut self,
last_event_id: EventId,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
match last_event_id.http_request_id {
Some(http_request_id) => {
let request_wise = self
.tx_router
.get_mut(&http_request_id)
.ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let (tx, rx) = channel;
request_wise.tx.tx = tx;
let index = last_event_id.index;
// sync messages after index
request_wise.tx.sync(index).await?;
Ok(StreamableHttpMessageReceiver {
http_request_id: Some(http_request_id),
inner: rx,
})
}
None => {
let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let (tx, rx) = channel;
self.common.tx = tx;
let index = last_event_id.index;
// sync messages after index
self.common.sync(index).await?;
Ok(StreamableHttpMessageReceiver {
http_request_id: None,
inner: rx,
})
}
}
}
}
enum SessionEvent {
ClientMessage {
pub trait SessionManager: Send + Sync + 'static {
type Error: std::error::Error + Send + 'static;
type Transport: crate::transport::Transport<RoleServer>;
/// Create a new session with the given id and configuration.
fn create_session(
&self,
) -> impl Future<Output = Result<(SessionId, Self::Transport), Self::Error>> + Send;
fn initialize_session(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
http_request_id: Option<HttpRequestId>,
},
EstablishRequestWiseChannel {
responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
},
CloseRequestWiseChannel {
id: HttpRequestId,
responder: oneshot::Sender<Result<(), SessionError>>,
},
Resume {
last_event_id: EventId,
responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
},
InitializeRequest {
request: ClientJsonRpcMessage,
responder: oneshot::Sender<Result<ServerJsonRpcMessage, SessionError>>,
},
Close,
}
#[derive(Debug, Clone)]
pub enum SessionQuitReason {
ServiceTerminated,
ClientTerminated,
ExpectInitializeRequest,
ExpectInitializeResponse,
Cancelled,
}
#[derive(Debug, Clone)]
pub struct SessionHandle {
id: SessionId,
// after all event_tx drop, inner task will be terminated
event_tx: Sender<SessionEvent>,
}
impl SessionHandle {
/// Get the session id
pub fn id(&self) -> &SessionId {
&self.id
}
/// Close the session
pub async fn close(&self) -> Result<(), SessionError> {
self.event_tx
.send(SessionEvent::Close)
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
Ok(())
}
/// Send a message to the session
pub async fn push_message(
) -> impl Future<Output = Result<ServerJsonRpcMessage, Self::Error>> + Send;
fn has_session(&self, id: &SessionId)
-> impl Future<Output = Result<bool, Self::Error>> + Send;
fn close_session(&self, id: &SessionId)
-> impl Future<Output = Result<(), Self::Error>> + Send;
fn create_stream(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
http_request_id: Option<HttpRequestId>,
) -> Result<(), SessionError> {
self.event_tx
.send(SessionEvent::ClientMessage {
message,
http_request_id,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
Ok(())
}
/// establish a channel for a http-request, the corresponded message from server will be
/// sent through this channel. The channel will be closed when the request is completed,
/// or you can close it manually by calling [`SessionHandle::close_request_wise_channel`].
pub async fn establish_request_wise_channel(
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send;
fn accept_message(
&self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::EstablishRequestWiseChannel { responder: tx })
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// close the http-request wise channel.
pub async fn close_request_wise_channel(
id: &SessionId,
message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn create_standalone_stream(
&self,
request_id: HttpRequestId,
) -> Result<(), SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::CloseRequestWiseChannel {
id: request_id,
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// Establish a common channel for general purpose messages.
pub async fn establish_common_channel(
id: &SessionId,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send;
fn resume(
&self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::Resume {
last_event_id: EventId {
http_request_id: None,
index: 0,
},
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// Resume streaming response by the last event id. This is suitable for both request wise and common channel.
pub async fn resume(
&self,
last_event_id: EventId,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::Resume {
last_event_id,
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// Send an initialize request to the session. And wait for the initialized response.
///
/// This is used to establish a session with the server.
pub async fn initialize(
&self,
request: ClientJsonRpcMessage,
) -> Result<ServerJsonRpcMessage, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::InitializeRequest {
request,
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
}
pub type SessionTransport = WorkerTransport<SessionWorker>;
impl Worker for SessionWorker {
type Error = SessionError;
type Role = RoleServer;
fn err_closed() -> Self::Error {
SessionError::TransportClosed
}
fn err_join(e: tokio::task::JoinError) -> Self::Error {
SessionError::TokioJoinError(e)
}
fn config(&self) -> crate::transport::worker::WorkerConfig {
crate::transport::worker::WorkerConfig {
name: Some(format!("streamable-http-session-{}", self.id)),
channel_buffer_capacity: self.session_config.channel_capacity,
}
}
#[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> {
enum InnerEvent {
FromHttpService(SessionEvent),
FromHandler(WorkerSendRequest<SessionWorker>),
}
// waiting for initialize request
let evt = self.event_rx.recv().await.ok_or_else(|| {
WorkerQuitReason::fatal("transport terminated", "get initialize request")
})?;
let SessionEvent::InitializeRequest { request, responder } = evt else {
return Err(WorkerQuitReason::fatal(
"unexpected message",
"get initialize request",
));
};
context.send_to_handler(request).await?;
let send_initialize_response = context.recv_from_handler().await?;
responder
.send(Ok(send_initialize_response.message))
.map_err(|_| {
WorkerQuitReason::fatal(
"failed to send initialize response to http service",
"send initialize response",
)
})?;
send_initialize_response
.responder
.send(Ok(()))
.map_err(|_| WorkerQuitReason::HandlerTerminated)?;
let ct = context.cancellation_token.clone();
loop {
let event = tokio::select! {
event = self.event_rx.recv() => {
if let Some(event) = event {
InnerEvent::FromHttpService(event)
} else {
return Err(WorkerQuitReason::fatal("session dropped", "waiting next session event"))
}
},
from_handler = context.recv_from_handler() => {
InnerEvent::FromHandler(from_handler?)
}
_ = ct.cancelled() => {
return Err(WorkerQuitReason::Cancelled)
}
};
match event {
InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => {
// catch response
match &message {
crate::model::JsonRpcMessage::Response(json_rpc_response) => {
let request_id = json_rpc_response.id.clone();
self.unregister_resource(&ResourceKey::McpRequestId(request_id));
}
crate::model::JsonRpcMessage::Error(json_rpc_error) => {
let request_id = json_rpc_error.id.clone();
self.unregister_resource(&ResourceKey::McpRequestId(request_id));
}
// unlikely happen
crate::model::JsonRpcMessage::BatchResponse(
json_rpc_batch_response_items,
) => {
for item in json_rpc_batch_response_items {
let request_id = match item {
crate::model::JsonRpcBatchResponseItem::Response(
json_rpc_response,
) => json_rpc_response.id.clone(),
crate::model::JsonRpcBatchResponseItem::Error(
json_rpc_error,
) => json_rpc_error.id.clone(),
};
self.unregister_resource(&ResourceKey::McpRequestId(request_id));
}
}
_ => {
// no need to unregister resource
}
}
let handle_result = self.handle_server_message(message).await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::ClientMessage {
message: json_rpc_message,
http_request_id,
}) => {
match &json_rpc_message {
crate::model::JsonRpcMessage::Request(request) => {
if let Some(http_request_id) = http_request_id {
self.register_request(request, http_request_id)
}
}
crate::model::JsonRpcMessage::Notification(notification) => {
self.catch_cancellation_notification(notification)
}
crate::model::JsonRpcMessage::BatchRequest(items) => {
for r in items {
match r {
crate::model::JsonRpcBatchRequestItem::Request(request) => {
if let Some(http_request_id) = http_request_id {
self.register_request(request, http_request_id)
}
}
crate::model::JsonRpcBatchRequestItem::Notification(
notification,
) => self.catch_cancellation_notification(notification),
}
}
}
_ => {}
}
context.send_to_handler(json_rpc_message).await?;
}
InnerEvent::FromHttpService(SessionEvent::EstablishRequestWiseChannel {
responder,
}) => {
let handle_result = self.establish_request_wise_channel().await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
id,
responder,
}) => {
let _handle_result = self.tx_router.remove(&id);
let _ = responder.send(Ok(()));
}
InnerEvent::FromHttpService(SessionEvent::Resume {
last_event_id,
responder,
}) => {
let handle_result = self.resume(last_event_id).await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::Close) => {
return Err(WorkerQuitReason::TransportClosed);
}
_ => {
// ignore
}
}
}
}
}
#[derive(Debug, Clone)]
pub struct SessionConfig {
channel_capacity: usize,
}
impl SessionConfig {
pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
}
}
}
/// Create a new session with the given id and configuration.
///
/// This function will return a pair of [`SessionHandle`] and [`SessionWorker`].
///
/// You can run the [`SessionWorker`] as a transport for mcp server. And use the [`SessionHandle`] operate the session.
pub fn create_session(
id: impl Into<SessionId>,
config: SessionConfig,
) -> (SessionHandle, SessionWorker) {
let id = id.into();
let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity);
let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity);
let common = CachedTx::new_common(common_tx);
tracing::info!(session_id = ?id, "create new session");
let handle = SessionHandle {
event_tx,
id: id.clone(),
};
let session_worker = SessionWorker {
next_http_request_id: 0,
id,
tx_router: HashMap::new(),
resource_router: HashMap::new(),
common,
event_rx,
session_config: config.clone(),
};
(handle, session_worker)
id: &SessionId,
last_event_id: String,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send;
}

View file

@ -0,0 +1,893 @@
use std::{
collections::{HashMap, HashSet, VecDeque},
num::ParseIntError,
sync::Arc,
time::Duration,
};
use futures::Stream;
use thiserror::Error;
use tokio::sync::{
mpsc::{Receiver, Sender},
oneshot,
};
use tokio_stream::wrappers::ReceiverStream;
use tracing::instrument;
use crate::{
RoleServer,
model::{
CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest,
JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
},
transport::{
WorkerTransport,
common::server_side_http::{SessionId, session_id},
worker::{Worker, WorkerContext, WorkerQuitReason, WorkerSendRequest},
},
};
#[derive(Debug, Default)]
pub struct LocalSessionManager {
pub sessions: tokio::sync::RwLock<HashMap<SessionId, LocalSessionHandle>>,
pub session_config: SessionConfig,
}
#[derive(Debug, Error)]
pub enum LocalSessionManagerError {
#[error("Session not found: {0}")]
SessionNotFound(SessionId),
#[error("Session error: {0}")]
SessionError(#[from] SessionError),
#[error("Invalid event id: {0}")]
InvalidEventId(#[from] EventIdParseError),
}
impl SessionManager for LocalSessionManager {
type Error = LocalSessionManagerError;
type Transport = WorkerTransport<LocalSessionWorker>;
async fn create_session(&self) -> Result<(SessionId, Self::Transport), Self::Error> {
let id = session_id();
let (handle, worker) = create_local_session(id.clone(), self.session_config.clone());
self.sessions.write().await.insert(id.clone(), handle);
Ok((id, WorkerTransport::spawn(worker)))
}
async fn initialize_session(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
) -> Result<ServerJsonRpcMessage, Self::Error> {
let sessions = self.sessions.read().await;
let handle = sessions
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
let response = handle.initialize(message).await?;
Ok(response)
}
async fn close_session(&self, id: &SessionId) -> Result<(), Self::Error> {
let mut sessions = self.sessions.write().await;
if let Some(handle) = sessions.remove(id) {
handle.close().await?;
}
Ok(())
}
async fn has_session(&self, id: &SessionId) -> Result<bool, Self::Error> {
let sessions = self.sessions.read().await;
Ok(sessions.contains_key(id))
}
async fn create_stream(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
let sessions = self.sessions.read().await;
let handle = sessions
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
let receiver = handle.establish_request_wise_channel().await?;
handle
.push_message(message, receiver.http_request_id)
.await?;
Ok(ReceiverStream::new(receiver.inner))
}
async fn create_standalone_stream(
&self,
id: &SessionId,
) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
let sessions = self.sessions.read().await;
let handle = sessions
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
let receiver = handle.establish_common_channel().await?;
Ok(ReceiverStream::new(receiver.inner))
}
async fn resume(
&self,
id: &SessionId,
last_event_id: String,
) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
let sessions = self.sessions.read().await;
let handle = sessions
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
let receiver = handle.resume(last_event_id.parse()?).await?;
Ok(ReceiverStream::new(receiver.inner))
}
async fn accept_message(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
) -> Result<(), Self::Error> {
let sessions = self.sessions.read().await;
let handle = sessions
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
handle.push_message(message, None).await?;
Ok(())
}
}
/// `<index>/request_id>`
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct EventId {
http_request_id: Option<HttpRequestId>,
index: usize,
}
impl std::fmt::Display for EventId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.index)?;
match &self.http_request_id {
Some(http_request_id) => write!(f, "/{http_request_id}"),
None => write!(f, ""),
}
}
}
#[derive(Debug, Clone, Error)]
pub enum EventIdParseError {
#[error("Invalid index: {0}")]
InvalidIndex(ParseIntError),
#[error("Invalid numeric request id: {0}")]
InvalidNumericRequestId(ParseIntError),
#[error("Missing request id type")]
InvalidRequestIdType,
#[error("Missing request id")]
MissingRequestId,
}
impl std::str::FromStr for EventId {
type Err = EventIdParseError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if let Some((index, request_id)) = s.split_once("/") {
let index = usize::from_str(index).map_err(EventIdParseError::InvalidIndex)?;
let request_id = u64::from_str(request_id).map_err(EventIdParseError::InvalidIndex)?;
Ok(EventId {
http_request_id: Some(request_id),
index,
})
} else {
let index = usize::from_str(s).map_err(EventIdParseError::InvalidIndex)?;
Ok(EventId {
http_request_id: None,
index,
})
}
}
}
use super::{ServerSseMessage, SessionManager};
struct CachedTx {
tx: Sender<ServerSseMessage>,
cache: VecDeque<ServerSseMessage>,
http_request_id: Option<HttpRequestId>,
capacity: usize,
}
impl CachedTx {
fn new(tx: Sender<ServerSseMessage>, http_request_id: Option<HttpRequestId>) -> Self {
Self {
cache: VecDeque::with_capacity(tx.capacity()),
capacity: tx.capacity(),
tx,
http_request_id,
}
}
fn new_common(tx: Sender<ServerSseMessage>) -> Self {
Self::new(tx, None)
}
async fn send(&mut self, message: ServerJsonRpcMessage) {
let index = self.cache.back().map_or(0, |m| {
m.event_id
.as_deref()
.unwrap_or_default()
.parse::<EventId>()
.expect("valid event id")
.index
+ 1
});
let event_id = EventId {
http_request_id: self.http_request_id,
index,
};
let message = ServerSseMessage {
event_id: Some(event_id.to_string()),
message: Arc::new(message),
};
if self.cache.len() >= self.capacity {
self.cache.pop_front();
self.cache.push_back(message.clone());
} else {
self.cache.push_back(message.clone());
}
let _ = self.tx.send(message).await.inspect_err(|e| {
let event_id = &e.0.event_id;
tracing::trace!(?event_id, "trying to send message in a closed session")
});
}
async fn sync(&mut self, index: usize) -> Result<(), SessionError> {
let Some(front) = self.cache.front() else {
return Ok(());
};
let front_event_id = front
.event_id
.as_deref()
.unwrap_or_default()
.parse::<EventId>()?;
let sync_index = index.saturating_sub(front_event_id.index);
if sync_index > self.cache.len() {
// invalid index
return Err(SessionError::InvalidEventId);
}
for message in self.cache.iter().skip(sync_index) {
let send_result = self.tx.send(message.clone()).await;
if send_result.is_err() {
let event_id: EventId = message.event_id.as_deref().unwrap_or_default().parse()?;
return Err(SessionError::ChannelClosed(Some(event_id.index as u64)));
}
}
Ok(())
}
}
struct HttpRequestWise {
resources: HashSet<ResourceKey>,
tx: CachedTx,
}
type HttpRequestId = u64;
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
enum ResourceKey {
McpRequestId(RequestId),
ProgressToken(ProgressToken),
}
pub struct LocalSessionWorker {
id: SessionId,
next_http_request_id: HttpRequestId,
tx_router: HashMap<HttpRequestId, HttpRequestWise>,
resource_router: HashMap<ResourceKey, HttpRequestId>,
common: CachedTx,
event_rx: Receiver<SessionEvent>,
session_config: SessionConfig,
}
impl LocalSessionWorker {
pub fn id(&self) -> &SessionId {
&self.id
}
}
#[derive(Debug, Error)]
pub enum SessionError {
#[error("Invalid request id: {0}")]
DuplicatedRequestId(HttpRequestId),
#[error("Channel closed: {0:?}")]
ChannelClosed(Option<HttpRequestId>),
#[error("Cannot parse event id: {0}")]
EventIdParseError(#[from] EventIdParseError),
#[error("Session service terminated")]
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 {
fn from(value: SessionError) -> Self {
match value {
SessionError::Io(io) => io,
_ => std::io::Error::new(std::io::ErrorKind::Other, format!("Session error: {value}")),
}
}
}
enum OutboundChannel {
RequestWise { id: HttpRequestId, close: bool },
Common,
}
pub struct StreamableHttpMessageReceiver {
pub http_request_id: Option<HttpRequestId>,
pub inner: Receiver<ServerSseMessage>,
}
impl LocalSessionWorker {
fn unregister_resource(&mut self, resource: &ResourceKey) {
if let Some(http_request_id) = self.resource_router.remove(resource) {
tracing::trace!(?resource, http_request_id, "unregister resource");
if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
channel.resources.remove(resource);
if channel.resources.is_empty() {
tracing::debug!(http_request_id, "close http request wise channel");
self.tx_router.remove(&http_request_id);
}
}
}
}
fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
tracing::trace!(?resource, http_request_id, "register resource");
if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
channel.resources.insert(resource.clone());
self.resource_router.insert(resource, http_request_id);
}
}
fn register_request(
&mut self,
request: &JsonRpcRequest<ClientRequest>,
http_request_id: HttpRequestId,
) {
use crate::model::GetMeta;
self.register_resource(
ResourceKey::McpRequestId(request.id.clone()),
http_request_id,
);
if let Some(progress_token) = request.request.get_meta().get_progress_token() {
self.register_resource(
ResourceKey::ProgressToken(progress_token.clone()),
http_request_id,
);
}
}
fn catch_cancellation_notification(
&mut self,
notification: &JsonRpcNotification<ClientNotification>,
) {
if let ClientNotification::CancelledNotification(n) = &notification.notification {
let request_id = n.params.request_id.clone();
let resource = ResourceKey::McpRequestId(request_id);
self.unregister_resource(&resource);
}
}
fn next_http_request_id(&mut self) -> HttpRequestId {
let id = self.next_http_request_id;
self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
id
}
async fn establish_request_wise_channel(
&mut self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let http_request_id = self.next_http_request_id();
let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
self.tx_router.insert(
http_request_id,
HttpRequestWise {
resources: Default::default(),
tx: CachedTx::new(tx, Some(http_request_id)),
},
);
tracing::debug!(http_request_id, "establish new request wise channel");
Ok(StreamableHttpMessageReceiver {
http_request_id: Some(http_request_id),
inner: rx,
})
}
fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel {
match &message {
ServerJsonRpcMessage::Request(_) => OutboundChannel::Common,
ServerJsonRpcMessage::Notification(JsonRpcNotification {
notification:
ServerNotification::ProgressNotification(Notification {
params: ProgressNotificationParam { progress_token, .. },
..
}),
..
}) => {
let id = self
.resource_router
.get(&ResourceKey::ProgressToken(progress_token.clone()));
if let Some(id) = id {
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::Notification(JsonRpcNotification {
notification:
ServerNotification::CancelledNotification(Notification {
params: CancelledNotificationParam { request_id, .. },
..
}),
..
}) => {
if let Some(id) = self
.resource_router
.get(&ResourceKey::McpRequestId(request_id.clone()))
{
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::Notification(_) => OutboundChannel::Common,
ServerJsonRpcMessage::Response(json_rpc_response) => {
if let Some(id) = self
.resource_router
.get(&ResourceKey::McpRequestId(json_rpc_response.id.clone()))
{
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::Error(json_rpc_error) => {
if let Some(id) = self
.resource_router
.get(&ResourceKey::McpRequestId(json_rpc_error.id.clone()))
{
OutboundChannel::RequestWise {
id: *id,
close: false,
}
} else {
OutboundChannel::Common
}
}
ServerJsonRpcMessage::BatchRequest(_) | ServerJsonRpcMessage::BatchResponse(_) => {
// the server side should never yield a batch request or response now
unreachable!("server side won't yield batch request or response")
}
}
}
async fn handle_server_message(
&mut self,
message: ServerJsonRpcMessage,
) -> Result<(), SessionError> {
let outbound_channel = self.resolve_outbound_channel(&message);
match outbound_channel {
OutboundChannel::RequestWise { id, close } => {
if let Some(request_wise) = self.tx_router.get_mut(&id) {
request_wise.tx.send(message).await;
if close {
self.tx_router.remove(&id);
}
} else {
return Err(SessionError::ChannelClosed(Some(id)));
}
}
OutboundChannel::Common => self.common.send(message).await,
}
Ok(())
}
async fn resume(
&mut self,
last_event_id: EventId,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
match last_event_id.http_request_id {
Some(http_request_id) => {
let request_wise = self
.tx_router
.get_mut(&http_request_id)
.ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let (tx, rx) = channel;
request_wise.tx.tx = tx;
let index = last_event_id.index;
// sync messages after index
request_wise.tx.sync(index).await?;
Ok(StreamableHttpMessageReceiver {
http_request_id: Some(http_request_id),
inner: rx,
})
}
None => {
let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let (tx, rx) = channel;
self.common.tx = tx;
let index = last_event_id.index;
// sync messages after index
self.common.sync(index).await?;
Ok(StreamableHttpMessageReceiver {
http_request_id: None,
inner: rx,
})
}
}
}
}
enum SessionEvent {
ClientMessage {
message: ClientJsonRpcMessage,
http_request_id: Option<HttpRequestId>,
},
EstablishRequestWiseChannel {
responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
},
CloseRequestWiseChannel {
id: HttpRequestId,
responder: oneshot::Sender<Result<(), SessionError>>,
},
Resume {
last_event_id: EventId,
responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
},
InitializeRequest {
request: ClientJsonRpcMessage,
responder: oneshot::Sender<Result<ServerJsonRpcMessage, SessionError>>,
},
Close,
}
#[derive(Debug, Clone)]
pub enum SessionQuitReason {
ServiceTerminated,
ClientTerminated,
ExpectInitializeRequest,
ExpectInitializeResponse,
Cancelled,
}
#[derive(Debug, Clone)]
pub struct LocalSessionHandle {
id: SessionId,
// after all event_tx drop, inner task will be terminated
event_tx: Sender<SessionEvent>,
}
impl LocalSessionHandle {
/// Get the session id
pub fn id(&self) -> &SessionId {
&self.id
}
/// Close the session
pub async fn close(&self) -> Result<(), SessionError> {
self.event_tx
.send(SessionEvent::Close)
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
Ok(())
}
/// Send a message to the session
pub async fn push_message(
&self,
message: ClientJsonRpcMessage,
http_request_id: Option<HttpRequestId>,
) -> Result<(), SessionError> {
self.event_tx
.send(SessionEvent::ClientMessage {
message,
http_request_id,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
Ok(())
}
/// establish a channel for a http-request, the corresponded message from server will be
/// sent through this channel. The channel will be closed when the request is completed,
/// or you can close it manually by calling [`LocalSessionHandle::close_request_wise_channel`].
pub async fn establish_request_wise_channel(
&self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::EstablishRequestWiseChannel { responder: tx })
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// close the http-request wise channel.
pub async fn close_request_wise_channel(
&self,
request_id: HttpRequestId,
) -> Result<(), SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::CloseRequestWiseChannel {
id: request_id,
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// Establish a common channel for general purpose messages.
pub async fn establish_common_channel(
&self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::Resume {
last_event_id: EventId {
http_request_id: None,
index: 0,
},
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// Resume streaming response by the last event id. This is suitable for both request wise and common channel.
pub async fn resume(
&self,
last_event_id: EventId,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::Resume {
last_event_id,
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
/// Send an initialize request to the session. And wait for the initialized response.
///
/// This is used to establish a session with the server.
pub async fn initialize(
&self,
request: ClientJsonRpcMessage,
) -> Result<ServerJsonRpcMessage, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::InitializeRequest {
request,
responder: tx,
})
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
.map_err(|_| SessionError::SessionServiceTerminated)?
}
}
pub type SessionTransport = WorkerTransport<LocalSessionWorker>;
impl Worker for LocalSessionWorker {
type Error = SessionError;
type Role = RoleServer;
fn err_closed() -> Self::Error {
SessionError::TransportClosed
}
fn err_join(e: tokio::task::JoinError) -> Self::Error {
SessionError::TokioJoinError(e)
}
fn config(&self) -> crate::transport::worker::WorkerConfig {
crate::transport::worker::WorkerConfig {
name: Some(format!("streamable-http-session-{}", self.id)),
channel_buffer_capacity: self.session_config.channel_capacity,
}
}
#[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> {
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")
})?;
let SessionEvent::InitializeRequest { request, responder } = evt else {
return Err(WorkerQuitReason::fatal(
"unexpected message",
"get initialize request",
));
};
context.send_to_handler(request).await?;
let send_initialize_response = context.recv_from_handler().await?;
responder
.send(Ok(send_initialize_response.message))
.map_err(|_| {
WorkerQuitReason::fatal(
"failed to send initialize response to http service",
"send initialize response",
)
})?;
send_initialize_response
.responder
.send(Ok(()))
.map_err(|_| WorkerQuitReason::HandlerTerminated)?;
let ct = context.cancellation_token.clone();
let keep_alive = self.session_config.keep_alive.unwrap_or(Duration::MAX);
loop {
let keep_alive_timeout = tokio::time::sleep(keep_alive);
let event = tokio::select! {
event = self.event_rx.recv() => {
if let Some(event) = event {
InnerEvent::FromHttpService(event)
} else {
return Err(WorkerQuitReason::fatal("session dropped", "waiting next session event"))
}
},
from_handler = context.recv_from_handler() => {
InnerEvent::FromHandler(from_handler?)
}
_ = ct.cancelled() => {
return Err(WorkerQuitReason::Cancelled)
}
_ = keep_alive_timeout => {
return Err(WorkerQuitReason::fatal("keep live timeout", "poll next session event"))
}
};
match event {
InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => {
// catch response
let to_unregister = match &message {
crate::model::JsonRpcMessage::Response(json_rpc_response) => {
let request_id = json_rpc_response.id.clone();
Some(ResourceKey::McpRequestId(request_id))
}
crate::model::JsonRpcMessage::Error(json_rpc_error) => {
let request_id = json_rpc_error.id.clone();
Some(ResourceKey::McpRequestId(request_id))
}
_ => {
None
// no need to unregister resource
}
};
let handle_result = self.handle_server_message(message).await;
let _ = responder.send(handle_result).inspect_err(|error| {
tracing::warn!(?error, "failed to send message to http service handler");
});
if let Some(to_unregister) = to_unregister {
self.unregister_resource(&to_unregister);
}
}
InnerEvent::FromHttpService(SessionEvent::ClientMessage {
message: json_rpc_message,
http_request_id,
}) => {
match &json_rpc_message {
crate::model::JsonRpcMessage::Request(request) => {
if let Some(http_request_id) = http_request_id {
self.register_request(request, http_request_id)
}
}
crate::model::JsonRpcMessage::Notification(notification) => {
self.catch_cancellation_notification(notification)
}
crate::model::JsonRpcMessage::BatchRequest(items) => {
for r in items {
match r {
crate::model::JsonRpcBatchRequestItem::Request(request) => {
if let Some(http_request_id) = http_request_id {
self.register_request(request, http_request_id)
}
}
crate::model::JsonRpcBatchRequestItem::Notification(
notification,
) => self.catch_cancellation_notification(notification),
}
}
}
_ => {}
}
context.send_to_handler(json_rpc_message).await?;
}
InnerEvent::FromHttpService(SessionEvent::EstablishRequestWiseChannel {
responder,
}) => {
let handle_result = self.establish_request_wise_channel().await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
id,
responder,
}) => {
let _handle_result = self.tx_router.remove(&id);
let _ = responder.send(Ok(()));
}
InnerEvent::FromHttpService(SessionEvent::Resume {
last_event_id,
responder,
}) => {
let handle_result = self.resume(last_event_id).await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::Close) => {
return Err(WorkerQuitReason::TransportClosed);
}
_ => {
// ignore
}
}
}
}
}
#[derive(Debug, Clone)]
pub struct SessionConfig {
/// the capacity of the channel for the session. Default is 16.
pub channel_capacity: usize,
/// if set, the session will be closed after this duration of inactivity.
pub keep_alive: Option<Duration>,
}
impl SessionConfig {
pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
keep_alive: None,
}
}
}
/// Create a new session with the given id and configuration.
///
/// This function will return a pair of [`LocalSessionHandle`] and [`LocalSessionWorker`].
///
/// You can run the [`LocalSessionWorker`] as a transport for mcp server. And use the [`LocalSessionHandle`] operate the session.
pub fn create_local_session(
id: impl Into<SessionId>,
config: SessionConfig,
) -> (LocalSessionHandle, LocalSessionWorker) {
let id = id.into();
let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity);
let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity);
let common = CachedTx::new_common(common_tx);
tracing::info!(session_id = ?id, "create new session");
let handle = LocalSessionHandle {
event_tx,
id: id.clone(),
};
let session_worker = LocalSessionWorker {
next_http_request_id: 0,
id,
tx_router: HashMap::new(),
resource_router: HashMap::new(),
common,
event_rx,
session_config: config.clone(),
};
(handle, session_worker)
}

View file

@ -0,0 +1,107 @@
use futures::Stream;
use thiserror::Error;
use super::{ServerSseMessage, SessionId, SessionManager};
use crate::{
RoleServer,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::Transport,
};
#[derive(Debug, Clone, Error)]
#[error("Session management is not supported")]
pub struct ErrorSessionManagementNotSupported;
#[derive(Debug, Clone, Default)]
pub struct NeverSessionManager {}
pub enum NeverTransport {}
impl Transport<RoleServer> for NeverTransport {
type Error = ErrorSessionManagementNotSupported;
fn send(
&mut self,
_item: ServerJsonRpcMessage,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn receive(&mut self) -> impl Future<Output = Option<ClientJsonRpcMessage>> {
futures::future::ready(None)
}
async fn close(&mut self) -> Result<(), Self::Error> {
Err(ErrorSessionManagementNotSupported)
}
}
impl SessionManager for NeverSessionManager {
type Error = ErrorSessionManagementNotSupported;
type Transport = NeverTransport;
fn create_session(
&self,
) -> impl Future<Output = Result<(SessionId, Self::Transport), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn initialize_session(
&self,
_id: &SessionId,
_message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<ServerJsonRpcMessage, Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn has_session(
&self,
_id: &SessionId,
) -> impl Future<Output = Result<bool, Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn close_session(
&self,
_id: &SessionId,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn create_stream(
&self,
_id: &SessionId,
_message: ClientJsonRpcMessage,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send {
futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
ErrorSessionManagementNotSupported,
))
}
fn create_standalone_stream(
&self,
_id: &SessionId,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send {
futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
ErrorSessionManagementNotSupported,
))
}
fn resume(
&self,
_id: &SessionId,
_last_event_id: String,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send {
futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
ErrorSessionManagementNotSupported,
))
}
fn accept_message(
&self,
_id: &SessionId,
_message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
}

View file

@ -0,0 +1,427 @@
use std::{convert::Infallible, fmt::Display, sync::Arc, time::Duration};
use bytes::Bytes;
use futures::{StreamExt, future::BoxFuture};
use http::{Method, Request, Response, header::ALLOW};
use http_body::Body;
use http_body_util::{BodyExt, Full, combinators::UnsyncBoxBody};
use tokio_stream::wrappers::ReceiverStream;
use super::session::SessionManager;
use crate::{
RoleServer,
model::ClientJsonRpcMessage,
serve_server,
service::serve_directly,
transport::{
OneshotTransport, TransportAdapterIdentity,
common::{
http_header::{
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
},
server_side_http::{
BoxResponse, ServerSseMessage, accepted_response, expect_json,
internal_error_response, sse_stream_response,
},
},
},
};
#[derive(Debug, Clone)]
pub struct StreamableHttpServerConfig {
/// The ping message duration for SSE connections.
pub sse_keep_alive: Option<Duration>,
/// If true, the server will create a session for each request and keep it alive.
pub stateful_mode: bool,
}
impl Default for StreamableHttpServerConfig {
fn default() -> Self {
Self {
sse_keep_alive: Some(Duration::from_secs(15)),
stateful_mode: true,
}
}
}
pub struct StreamableHttpService<S, M = super::session::local::LocalSessionManager> {
pub config: StreamableHttpServerConfig,
session_manager: Arc<M>,
service_factory: Arc<dyn Fn() -> S + Send + Sync>,
}
impl<S, M> Clone for StreamableHttpService<S, M> {
fn clone(&self) -> Self {
Self {
config: self.config.clone(),
session_manager: self.session_manager.clone(),
service_factory: self.service_factory.clone(),
}
}
}
impl<RequestBody, S, M> tower_service::Service<Request<RequestBody>> for StreamableHttpService<S, M>
where
RequestBody: Body + Send + 'static,
S: crate::Service<RoleServer>,
M: SessionManager,
RequestBody::Error: Display,
RequestBody::Data: Send + 'static,
{
type Response = BoxResponse;
type Error = Infallible;
type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn call(&mut self, req: http::Request<RequestBody>) -> Self::Future {
let service = self.clone();
Box::pin(async move {
let response = service.handle(req).await;
Ok(response)
})
}
fn poll_ready(
&mut self,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
}
impl<S, M> StreamableHttpService<S, M>
where
S: crate::Service<RoleServer> + Send + 'static,
M: SessionManager,
{
pub fn new(
service_factory: impl Fn() -> S + Send + Sync + 'static,
session_manager: Arc<M>,
config: StreamableHttpServerConfig,
) -> Self {
Self {
config,
session_manager,
service_factory: Arc::new(service_factory),
}
}
fn get_service(&self) -> S {
(self.service_factory)()
}
pub async fn handle<B>(&self, request: Request<B>) -> Response<UnsyncBoxBody<Bytes, Infallible>>
where
B: Body + Send + 'static,
B::Error: Display,
{
let method = request.method().clone();
let result = match method {
Method::GET => self.handle_get(request).await,
Method::POST => self.handle_post(request).await,
Method::DELETE => self.handle_delete(request).await,
_ => {
// Handle other methods or return an error
let response = Response::builder()
.status(http::StatusCode::METHOD_NOT_ALLOWED)
.header(ALLOW, "GET, POST, DELETE")
.body(Full::new(Bytes::from("Method Not Allowed")).boxed_unsync())
.expect("valid response");
return response;
}
};
match result {
Ok(response) => response,
Err(response) => response,
}
}
async fn handle_get<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
where
B: Body + Send + 'static,
B::Error: Display,
{
// check accept header
if !request
.headers()
.get(http::header::ACCEPT)
.and_then(|header| header.to_str().ok())
.is_some_and(|header| header.contains(EVENT_STREAM_MIME_TYPE))
{
return Ok(Response::builder()
.status(http::StatusCode::NOT_ACCEPTABLE)
.body(
Full::new(Bytes::from(
"Not Acceptable: Client must accept text/event-stream",
))
.boxed_unsync(),
)
.expect("valid response"));
}
// check session id
let session_id = request
.headers()
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned().into());
let Some(session_id) = session_id else {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session ID is required")).boxed_unsync())
.expect("valid response"));
};
// check if session exists
let has_session = self
.session_manager
.has_session(&session_id)
.await
.map_err(internal_error_response("check session"))?;
if !has_session {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session not found")).boxed_unsync())
.expect("valid response"));
}
// check if last event id is provided
let last_event_id = request
.headers()
.get(HEADER_LAST_EVENT_ID)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned());
if let Some(last_event_id) = last_event_id {
// check if session has this event id
let stream = self
.session_manager
.resume(&session_id, last_event_id)
.await
.map_err(internal_error_response("resume session"))?;
Ok(sse_stream_response(stream, self.config.sse_keep_alive))
} else {
// create standalone stream
let stream = self
.session_manager
.create_standalone_stream(&session_id)
.await
.map_err(internal_error_response("create standalone stream"))?;
Ok(sse_stream_response(stream, self.config.sse_keep_alive))
}
}
async fn handle_post<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
where
B: Body + Send + 'static,
B::Error: Display,
{
// check accept header
if !request
.headers()
.get(http::header::ACCEPT)
.and_then(|header| header.to_str().ok())
.is_some_and(|header| {
header.contains(JSON_MIME_TYPE) && header.contains(EVENT_STREAM_MIME_TYPE)
})
{
return Ok(Response::builder()
.status(http::StatusCode::NOT_ACCEPTABLE)
.body(Full::new(Bytes::from("Not Acceptable: Client must accept both application/json and text/event-stream")).boxed_unsync())
.expect("valid response"));
}
// check content type
if !request
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|header| header.to_str().ok())
.is_some_and(|header| header.starts_with(JSON_MIME_TYPE))
{
return Ok(Response::builder()
.status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
.body(
Full::new(Bytes::from(
"Unsupported Media Type: Content-Type must be application/json",
))
.boxed_unsync(),
)
.expect("valid response"));
}
// json deserialize request body
let (part, body) = request.into_parts();
let message = match expect_json(body).await {
Ok(message) => message,
Err(response) => return Ok(response),
};
if self.config.stateful_mode {
// do we have a session id?
let session_id = part
.headers
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok());
if let Some(session_id) = session_id {
let session_id = session_id.to_owned().into();
let has_session = self
.session_manager
.has_session(&session_id)
.await
.map_err(internal_error_response("check session"))?;
if !has_session {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(
Full::new(Bytes::from("Unauthorized: Session not found"))
.boxed_unsync(),
)
.expect("valid response"));
}
match message {
ClientJsonRpcMessage::Request(_) => {
let stream = self
.session_manager
.create_stream(&session_id, message)
.await
.map_err(internal_error_response("get session"))?;
Ok(sse_stream_response(stream, self.config.sse_keep_alive))
}
ClientJsonRpcMessage::Notification(_)
| ClientJsonRpcMessage::Response(_)
| ClientJsonRpcMessage::Error(_) => {
// handle notification
self.session_manager
.accept_message(&session_id, message)
.await
.map_err(internal_error_response("accept message"))?;
Ok(accepted_response())
}
_ => Ok(Response::builder()
.status(http::StatusCode::NOT_IMPLEMENTED)
.body(
Full::new(Bytes::from("Batch requests are not supported yet"))
.boxed_unsync(),
)
.expect("valid response")),
}
} else {
let (session_id, transport) = self
.session_manager
.create_session()
.await
.map_err(internal_error_response("create session"))?;
let service = self.get_service();
// spawn a task to serve the session
tokio::spawn({
let session_manager = self.session_manager.clone();
let session_id = session_id.clone();
async move {
let service = serve_server::<S, M::Transport, _, TransportAdapterIdentity>(
service, transport,
)
.await;
match service {
Ok(service) => {
// on service created
let _ = service.waiting().await;
}
Err(e) => {
tracing::error!("Failed to create service: {e}");
}
}
let _ = session_manager
.close_session(&session_id)
.await
.inspect_err(|e| {
tracing::error!("Failed to close session {session_id}: {e}");
});
}
});
// get initialize response
let response = self
.session_manager
.initialize_session(&session_id, message)
.await
.map_err(internal_error_response("create stream"))?;
let mut response = sse_stream_response(
futures::stream::once({
async move {
ServerSseMessage {
event_id: None,
message: response.into(),
}
}
}),
self.config.sse_keep_alive,
);
response.headers_mut().insert(
HEADER_SESSION_ID,
session_id
.parse()
.map_err(internal_error_response("create session id header"))?,
);
Ok(response)
}
} else {
let service = self.get_service();
match message {
ClientJsonRpcMessage::Request(request) => {
let (transport, receiver) =
OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
let service = serve_directly(service, transport, None);
tokio::spawn(async move {
// on service created
let _ = service.waiting().await;
});
Ok(sse_stream_response(
ReceiverStream::new(receiver).map(|message| {
tracing::info!(?message);
ServerSseMessage {
event_id: None,
message: message.into(),
}
}),
self.config.sse_keep_alive,
))
}
ClientJsonRpcMessage::Notification(notification) => {
service
.handle_notification(notification.notification)
.await
.map_err(internal_error_response("handle notification"))?;
Ok(accepted_response())
}
ClientJsonRpcMessage::Response(_json_rpc_response) => Ok(accepted_response()),
ClientJsonRpcMessage::Error(_json_rpc_error) => Ok(accepted_response()),
_ => Ok(Response::builder()
.status(http::StatusCode::NOT_IMPLEMENTED)
.body(
Full::new(Bytes::from("Batch requests are not supported yet"))
.boxed_unsync(),
)
.expect("valid response")),
}
}
}
async fn handle_delete<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
where
B: Body + Send + 'static,
B::Error: Display,
{
// check session id
let session_id = request
.headers()
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned().into());
let Some(session_id) = session_id else {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session ID is required")).boxed_unsync())
.expect("valid response"));
};
// close session
self.session_manager
.close_session(&session_id)
.await
.map_err(internal_error_response("close session"))?;
Ok(accepted_response())
}
}

View file

@ -2,10 +2,14 @@ use rmcp::{
ServiceExt,
service::QuitReason,
transport::{
ConfigureCommandExt, SseServer, StreamableHttpClientTransport, TokioChildProcess,
streamable_http_server::axum::StreamableHttpServer,
ConfigureCommandExt, SseServer, StreamableHttpClientTransport, StreamableHttpServerConfig,
TokioChildProcess,
streamable_http_server::{
session::local::LocalSessionManager, tower::StreamableHttpService,
},
},
};
use tokio_util::sync::CancellationToken;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
mod common;
use common::calculator::Calculator;
@ -90,10 +94,26 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
.wait()
.await?;
let ct = StreamableHttpServer::serve(STREAMABLE_HTTP_BIND_ADDRESS.parse()?)
.await?
.with_service(Calculator::default);
let service: StreamableHttpService<Calculator, LocalSessionManager> =
StreamableHttpService::new(
Calculator::default,
Default::default(),
StreamableHttpServerConfig {
stateful_mode: true,
sse_keep_alive: None,
},
);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind(STREAMABLE_HTTP_BIND_ADDRESS).await?;
let ct = CancellationToken::new();
let handle = tokio::spawn({
let ct = ct.clone();
async move {
let _ = axum::serve(tcp_listener, router)
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
.await;
}
});
let exit_status = tokio::process::Command::new("node")
.arg("tests/test_with_js/streamable_client.js")
.spawn()?
@ -101,6 +121,7 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
.await?;
assert!(exit_status.success());
ct.cancel();
handle.await?;
Ok(())
}

View file

@ -1,7 +1,7 @@
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
const transport = new StreamableHTTPClientTransport(new URL(`http://127.0.0.1:8001/`));
const transport = new StreamableHTTPClientTransport(new URL(`http://127.0.0.1:8001/mcp/`));
const client = new Client(
{

View file

@ -7,8 +7,20 @@ edition = "2024"
publish = false
[dependencies]
rmcp= { path = "../../crates/rmcp", features = ["server", "transport-sse-server", "transport-io", "transport-streamable-http-server", "auth"] }
tokio = { version = "1", features = ["macros", "rt", "rt-multi-thread", "io-std", "signal"] }
rmcp = { path = "../../crates/rmcp", features = [
"server",
"transport-sse-server",
"transport-io",
"transport-streamable-http-server",
"auth",
] }
tokio = { version = "1", features = [
"macros",
"rt",
"rt-multi-thread",
"io-std",
"signal",
] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
anyhow = "1.0"
@ -26,9 +38,10 @@ reqwest = { version = "0.12", features = ["json"] }
chrono = "0.4"
uuid = { version = "1.6", features = ["v4", "serde"] }
serde_urlencoded = "0.7"
askama = { version = "0.14"}
askama = { version = "0.14" }
tower-http = { version = "0.6", features = ["cors"] }
hyper = { version = "1" }
hyper-util = { version = "0", features = ["server"] }
[dev-dependencies]
tokio-stream = { version = "0.1" }
@ -60,4 +73,8 @@ path = "src/complex_auth_sse.rs"
[[example]]
name = "servers_simple_auth_sse"
path = "src/simple_auth_sse.rs"
path = "src/simple_auth_sse.rs"
[[example]]
name = "counter_hyper_streamable_http"
path = "src/counter_hyper_streamable_http.rs"

View file

@ -38,12 +38,19 @@ A minimal server example using stdio transport.
### Counter Streamable HTTP Server (`counter_streamhttp.rs`)
A server using streamable HTTP transport for MCP communication.
A server using streamable HTTP transport for MCP communication, with axum.
- Runs on HTTP with streaming capabilities
- Provides counter tools via HTTP streaming
- Demonstrates streamable HTTP transport configuration
### Counter Streamable HTTP Server with Hyper (`counter_hyper_streamable_http.rs`)
A server using streamable HTTP transport for MCP communication, with hyper.
- Runs on HTTP with streaming capabilities
- Provides counter tools via HTTP streaming
- Demonstrates streamable HTTP transport configuration
### Complex OAuth SSE Server (`complex_auth_sse.rs`)
A comprehensive example demonstrating OAuth 2.0 integration with MCP servers.

Some files were not shown because too many files have changed in this diff Show more