refactor: provide http server as tower service (#228)
* refactor: streamable http server as tower service
This commit is contained in:
parent
915bc3fb37
commit
209be7bdb8
22 changed files with 1828 additions and 1180 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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>,
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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(()))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
140
crates/rmcp/src/transport/common/server_side_http.rs
Normal file
140
crates/rmcp/src/transport/common/server_side_http.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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(())
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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) = ¬ification.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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) = ¬ification.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)
|
||||
}
|
||||
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
427
crates/rmcp/src/transport/streamable_http_server/tower.rs
Normal file
427
crates/rmcp/src/transport/streamable_http_server/tower.rs
Normal 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())
|
||||
}
|
||||
}
|
||||
|
|
@ -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(())
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Reference in a new issue