feat!: remove SSE transport support (#562)
SSE transport has been removed from the MCP specification in favor of streamable HTTP. This removes all SSE-specific transport code: - Remove `transport-sse-client` and `transport-sse-server` features - Remove `SseClientTransport` and `SseServer` types - Remove SSE-specific examples (`counter_sse`, `counter_sse_directly`) - Migrate auth examples from SSE to streamable HTTP - Update tests to remove SSE transport usage - Update documentation BREAKING CHANGE: The following have been removed: - `transport-sse-client` feature - `transport-sse-client-reqwest` feature - `transport-sse-server` feature - `SseClientTransport` type - `SseServer` type - `sse_client` and `sse_server` modules Users should migrate to streamable HTTP transport which provides equivalent functionality. See `StreamableHttpClientTransport` and `StreamableHttpService` for the replacement APIs. Ref: https://github.com/modelcontextprotocol/rust-sdk/pull/561#issuecomment-3576551699
This commit is contained in:
adam jones
GitHub
parent
3dee024325
commit
eb5a7f7408
9 changed files with 164 additions and 1688 deletions
|
|
@ -2,13 +2,12 @@
|
|||
//! The transport type must implemented [`Transport`] trait, which allow it send message concurrently and receive message sequentially.
|
||||
//！
|
||||
//! ## Standard Transport Types
|
||||
//! There are 3 pairs of standard transport types:
|
||||
//! There are 2 pairs of standard transport types:
|
||||
//!
|
||||
//! | transport | client | server |
|
||||
//! |:-: |:-: |:-: |
|
||||
//! | std IO | [`child_process::TokioChildProcess`] | [`io::stdio`] |
|
||||
//! | streamable http | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::StreamableHttpService`] |
|
||||
//! | sse | [`sse_client::SseClientTransport`] | [`sse_server::SseServer`] |
|
||||
//!
|
||||
//！## Helper Transport Types
|
||||
//! Thers are several helper transport types that can help you to create transport quickly.
|
||||
|
|
@ -95,20 +94,6 @@ pub mod io;
|
|||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-io")))]
|
||||
pub use io::stdio;
|
||||
|
||||
#[cfg(feature = "transport-sse-client")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-client")))]
|
||||
pub mod sse_client;
|
||||
#[cfg(feature = "transport-sse-client")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-client")))]
|
||||
pub use sse_client::SseClientTransport;
|
||||
|
||||
#[cfg(feature = "transport-sse-server")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-server")))]
|
||||
pub mod sse_server;
|
||||
#[cfg(feature = "transport-sse-server")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-server")))]
|
||||
pub use sse_server::SseServer;
|
||||
|
||||
#[cfg(feature = "auth")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
|
||||
pub mod auth;
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ impl CredentialStore for InMemoryCredentialStore {
|
|||
}
|
||||
}
|
||||
|
||||
/// sse client with oauth2 authorization
|
||||
/// HTTP client with OAuth 2.0 authorization
|
||||
#[derive(Clone)]
|
||||
pub struct AuthClient<C> {
|
||||
pub http_client: C,
|
||||
|
|
@ -91,7 +91,7 @@ impl<C: std::fmt::Debug> std::fmt::Debug for AuthClient<C> {
|
|||
}
|
||||
|
||||
impl<C> AuthClient<C> {
|
||||
/// create new authorized sse client
|
||||
/// Create a new authorized HTTP client
|
||||
pub fn new(http_client: C, auth_manager: AuthorizationManager) -> Self {
|
||||
Self {
|
||||
http_client,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,4 @@
|
|||
#[cfg(any(
|
||||
feature = "transport-streamable-http-server",
|
||||
feature = "transport-sse-server"
|
||||
))]
|
||||
#[cfg(feature = "transport-streamable-http-server")]
|
||||
pub mod server_side_http;
|
||||
|
||||
pub mod http_header;
|
||||
|
|
@ -10,6 +7,9 @@ pub mod http_header;
|
|||
#[cfg_attr(docsrs, doc(cfg(feature = "reqwest")))]
|
||||
mod reqwest;
|
||||
|
||||
// Note: This module provides SSE stream parsing and auto-reconnect utilities.
|
||||
// It's used by the streamable HTTP client (which receives SSE-formatted responses),
|
||||
// not the removed SSE transport. The name is historical.
|
||||
#[cfg(feature = "client-side-sse")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "client-side-sse")))]
|
||||
pub mod client_side_sse;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,3 @@
|
|||
#[cfg(feature = "transport-streamable-http-client")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client")))]
|
||||
mod streamable_http_client;
|
||||
|
||||
#[cfg(feature = "transport-sse-client")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-client")))]
|
||||
mod sse_client;
|
||||
|
|
|
|||
|
|
@ -1,45 +0,0 @@
|
|||
use http::Uri;
|
||||
|
||||
use crate::transport::{
|
||||
auth::AuthClient,
|
||||
sse_client::{SseClient, SseTransportError},
|
||||
};
|
||||
impl<C> SseClient for AuthClient<C>
|
||||
where
|
||||
C: SseClient,
|
||||
{
|
||||
type Error = SseTransportError<C::Error>;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
uri: Uri,
|
||||
message: crate::model::ClientJsonRpcMessage,
|
||||
mut auth_token: Option<String>,
|
||||
) -> Result<(), SseTransportError<Self::Error>> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.post_message(uri, message, auth_token)
|
||||
.await
|
||||
.map_err(SseTransportError::Client)
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
uri: Uri,
|
||||
last_event_id: Option<String>,
|
||||
mut auth_token: Option<String>,
|
||||
) -> Result<
|
||||
crate::transport::common::client_side_sse::BoxedSseResponse,
|
||||
SseTransportError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.get_stream(uri, last_event_id, auth_token)
|
||||
.await
|
||||
.map_err(SseTransportError::Client)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,3 @@
|
|||
#[cfg(feature = "transport-streamable-http-client-reqwest")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client-reqwest")))]
|
||||
mod streamable_http_client;
|
||||
|
||||
#[cfg(feature = "transport-sse-client-reqwest")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-client-reqwest")))]
|
||||
mod sse_client;
|
||||
|
|
|
|||
|
|
@ -1,118 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use futures::StreamExt;
|
||||
use http::Uri;
|
||||
use reqwest::header::ACCEPT;
|
||||
use sse_stream::SseStream;
|
||||
|
||||
use crate::transport::{
|
||||
SseClientTransport,
|
||||
common::http_header::{EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID},
|
||||
sse_client::{SseClient, SseClientConfig, SseTransportError},
|
||||
};
|
||||
|
||||
impl From<reqwest::Error> for SseTransportError<reqwest::Error> {
|
||||
fn from(e: reqwest::Error) -> Self {
|
||||
SseTransportError::Client(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl SseClient for reqwest::Client {
|
||||
type Error = reqwest::Error;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
uri: Uri,
|
||||
message: crate::model::ClientJsonRpcMessage,
|
||||
auth_token: Option<String>,
|
||||
) -> Result<(), SseTransportError<Self::Error>> {
|
||||
let mut request_builder = self.post(uri.to_string()).json(&message);
|
||||
if let Some(auth_header) = auth_token {
|
||||
request_builder = request_builder.bearer_auth(auth_header);
|
||||
}
|
||||
request_builder
|
||||
.send()
|
||||
.await
|
||||
.and_then(|resp| resp.error_for_status())
|
||||
.map_err(SseTransportError::from)
|
||||
.map(drop)
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
uri: Uri,
|
||||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
) -> Result<
|
||||
crate::transport::common::client_side_sse::BoxedSseResponse,
|
||||
SseTransportError<Self::Error>,
|
||||
> {
|
||||
let mut request_builder = self
|
||||
.get(uri.to_string())
|
||||
.header(ACCEPT, EVENT_STREAM_MIME_TYPE);
|
||||
if let Some(auth_header) = auth_token {
|
||||
request_builder = request_builder.bearer_auth(auth_header);
|
||||
}
|
||||
if let Some(last_event_id) = last_event_id {
|
||||
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
|
||||
}
|
||||
let response = request_builder.send().await?;
|
||||
let response = response.error_for_status()?;
|
||||
match response.headers().get(reqwest::header::CONTENT_TYPE) {
|
||||
Some(ct) => {
|
||||
if !ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) {
|
||||
return Err(SseTransportError::UnexpectedContentType(Some(
|
||||
String::from_utf8_lossy(ct.as_bytes()).to_string(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
None => {
|
||||
return Err(SseTransportError::UnexpectedContentType(None));
|
||||
}
|
||||
}
|
||||
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
|
||||
Ok(event_stream)
|
||||
}
|
||||
}
|
||||
|
||||
impl SseClientTransport<reqwest::Client> {
|
||||
/// Creates a new transport using reqwest with the specified SSE endpoint.
|
||||
///
|
||||
/// This is a convenience method that creates a transport using the default
|
||||
/// reqwest client. This method is only available when the
|
||||
/// `transport-sse-client-reqwest` feature is enabled.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `uri` - The SSE endpoint to connect to
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use rmcp::transport::SseClientTransport;
|
||||
///
|
||||
/// // Enable the reqwest feature in Cargo.toml:
|
||||
/// // rmcp = { version = "0.5", features = ["transport-sse-client-reqwest"] }
|
||||
///
|
||||
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let transport = SseClientTransport::start("http://localhost:8000/sse").await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
///
|
||||
/// # Feature requirement
|
||||
///
|
||||
/// This method requires the `transport-sse-client-reqwest` feature.
|
||||
pub async fn start(
|
||||
uri: impl Into<Arc<str>>,
|
||||
) -> Result<Self, SseTransportError<reqwest::Error>> {
|
||||
SseClientTransport::start_with_client(
|
||||
reqwest::Client::default(),
|
||||
SseClientConfig {
|
||||
sse_endpoint: uri.into(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
@ -1,458 +0,0 @@
|
|||
//! Reference: <https://html.spec.whatwg.org/multipage/server-sent-events.html>
|
||||
use std::{
|
||||
pin::Pin,
|
||||
sync::{Arc, RwLock},
|
||||
};
|
||||
|
||||
use futures::{StreamExt, future::BoxFuture};
|
||||
use http::Uri;
|
||||
use sse_stream::{Error as SseError, Sse};
|
||||
use thiserror::Error;
|
||||
|
||||
use super::{
|
||||
Transport,
|
||||
common::client_side_sse::{BoxedSseResponse, SseRetryPolicy, SseStreamReconnect},
|
||||
};
|
||||
use crate::{
|
||||
RoleClient,
|
||||
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
|
||||
transport::common::client_side_sse::SseAutoReconnectStream,
|
||||
};
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum SseTransportError<E: std::error::Error + Send + Sync + 'static> {
|
||||
#[error("SSE error: {0}")]
|
||||
Sse(#[from] SseError),
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
#[error("Client error: {0}")]
|
||||
Client(E),
|
||||
#[error("unexpected end of stream")]
|
||||
UnexpectedEndOfStream,
|
||||
#[error("Unexpected content type: {0:?}")]
|
||||
UnexpectedContentType(Option<String>),
|
||||
#[cfg(feature = "auth")]
|
||||
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
|
||||
#[error("Auth error: {0}")]
|
||||
Auth(#[from] crate::transport::auth::AuthError),
|
||||
#[error("Invalid uri: {0}")]
|
||||
InvalidUri(#[from] http::uri::InvalidUri),
|
||||
#[error("Invalid uri parts: {0}")]
|
||||
InvalidUriParts(#[from] http::uri::InvalidUriParts),
|
||||
}
|
||||
|
||||
pub trait SseClient: Clone + Send + Sync + 'static {
|
||||
type Error: std::error::Error + Send + Sync + 'static;
|
||||
fn post_message(
|
||||
&self,
|
||||
uri: Uri,
|
||||
message: ClientJsonRpcMessage,
|
||||
auth_token: Option<String>,
|
||||
) -> impl Future<Output = Result<(), SseTransportError<Self::Error>>> + Send + '_;
|
||||
fn get_stream(
|
||||
&self,
|
||||
uri: Uri,
|
||||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
) -> impl Future<Output = Result<BoxedSseResponse, SseTransportError<Self::Error>>> + Send + '_;
|
||||
}
|
||||
|
||||
/// Helper that refreshes the POST endpoint whenever the server emits
|
||||
/// control frames during SSE reconnect; used together with
|
||||
/// [`SseAutoReconnectStream`].
|
||||
struct SseClientReconnect<C> {
|
||||
pub client: C,
|
||||
pub uri: Uri,
|
||||
pub message_endpoint: Arc<RwLock<Uri>>,
|
||||
}
|
||||
|
||||
impl<C: SseClient> SseStreamReconnect for SseClientReconnect<C> {
|
||||
type Error = SseTransportError<C::Error>;
|
||||
type Future = BoxFuture<'static, Result<BoxedSseResponse, Self::Error>>;
|
||||
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future {
|
||||
let client = self.client.clone();
|
||||
let uri = self.uri.clone();
|
||||
let last_event_id = last_event_id.map(|s| s.to_owned());
|
||||
Box::pin(async move { client.get_stream(uri, last_event_id, None).await })
|
||||
}
|
||||
|
||||
fn handle_control_event(&mut self, event: &Sse) -> Result<(), Self::Error> {
|
||||
if event.event.as_deref() != Some("endpoint") {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(data) = event.data.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
// Servers typically resend the message POST endpoint (often with a new
|
||||
// sessionId) when a stream reconnects. Reuse `message_endpoint` helper
|
||||
// to resolve it and update the shared URI.
|
||||
let new_endpoint = message_endpoint(self.uri.clone(), data.clone())
|
||||
.map_err(SseTransportError::InvalidUri)?;
|
||||
*self
|
||||
.message_endpoint
|
||||
.write()
|
||||
.expect("message endpoint lock poisoned") = new_endpoint;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_stream_error(
|
||||
&mut self,
|
||||
error: &(dyn std::error::Error + 'static),
|
||||
last_event_id: Option<&str>,
|
||||
) {
|
||||
tracing::warn!(
|
||||
uri = %self.uri,
|
||||
last_event_id = last_event_id.unwrap_or(""),
|
||||
"sse stream error: {error}"
|
||||
);
|
||||
}
|
||||
}
|
||||
type ServerMessageStream<C> = Pin<Box<SseAutoReconnectStream<SseClientReconnect<C>>>>;
|
||||
|
||||
/// A client-agnostic SSE transport for RMCP that supports Server-Sent Events.
|
||||
///
|
||||
/// This transport allows you to choose your preferred HTTP client implementation
|
||||
/// by implementing the [`SseClient`] trait. The transport handles SSE streaming
|
||||
/// and automatic reconnection.
|
||||
///
|
||||
/// # Usage
|
||||
///
|
||||
/// ## Using reqwest
|
||||
///
|
||||
/// ```rust,ignore
|
||||
/// use rmcp::transport::SseClientTransport;
|
||||
///
|
||||
/// // Enable the reqwest feature in Cargo.toml:
|
||||
/// // rmcp = { version = "0.5", features = ["transport-sse-client-reqwest"] }
|
||||
///
|
||||
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let transport = SseClientTransport::start("http://localhost:8000/sse").await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
///
|
||||
/// ## Using a custom HTTP client
|
||||
///
|
||||
/// ```rust,ignore
|
||||
/// use rmcp::transport::sse_client::{SseClient, SseClientTransport, SseClientConfig};
|
||||
/// use std::sync::Arc;
|
||||
/// use futures::stream::BoxStream;
|
||||
/// use rmcp::model::ClientJsonRpcMessage;
|
||||
/// use sse_stream::{Sse, Error as SseError};
|
||||
/// use http::Uri;
|
||||
///
|
||||
/// #[derive(Clone)]
|
||||
/// struct MyHttpClient;
|
||||
///
|
||||
/// #[derive(Debug, thiserror::Error)]
|
||||
/// struct MyError;
|
||||
///
|
||||
/// impl std::fmt::Display for MyError {
|
||||
/// fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
/// write!(f, "MyError")
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// impl SseClient for MyHttpClient {
|
||||
/// type Error = MyError;
|
||||
///
|
||||
/// async fn post_message(
|
||||
/// &self,
|
||||
/// _uri: Uri,
|
||||
/// _message: ClientJsonRpcMessage,
|
||||
/// _auth_token: Option<String>,
|
||||
/// ) -> Result<(), rmcp::transport::sse_client::SseTransportError<Self::Error>> {
|
||||
/// todo!()
|
||||
/// }
|
||||
///
|
||||
/// async fn get_stream(
|
||||
/// &self,
|
||||
/// _uri: Uri,
|
||||
/// _last_event_id: Option<String>,
|
||||
/// _auth_token: Option<String>,
|
||||
/// ) -> Result<BoxStream<'static, Result<Sse, SseError>>, rmcp::transport::sse_client::SseTransportError<Self::Error>> {
|
||||
/// todo!()
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let config = SseClientConfig {
|
||||
/// sse_endpoint: "http://localhost:8000/sse".into(),
|
||||
/// ..Default::default()
|
||||
/// };
|
||||
/// let transport = SseClientTransport::start_with_client(MyHttpClient, config).await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
///
|
||||
/// # Feature Flags
|
||||
///
|
||||
/// - `transport-sse-client`: Base feature providing the generic transport infrastructure
|
||||
/// - `transport-sse-client-reqwest`: Includes reqwest HTTP client support with convenience methods
|
||||
pub struct SseClientTransport<C: SseClient> {
|
||||
client: C,
|
||||
config: SseClientConfig,
|
||||
/// Current POST endpoint; refreshed when the server sends new endpoint
|
||||
/// control frames.
|
||||
message_endpoint: Arc<RwLock<Uri>>,
|
||||
stream: Option<ServerMessageStream<C>>,
|
||||
}
|
||||
|
||||
impl<C: SseClient> Transport<RoleClient> for SseClientTransport<C> {
|
||||
type Error = SseTransportError<C::Error>;
|
||||
async fn receive(&mut self) -> Option<ServerJsonRpcMessage> {
|
||||
self.stream.as_mut()?.next().await?.ok()
|
||||
}
|
||||
fn send(
|
||||
&mut self,
|
||||
item: crate::service::TxJsonRpcMessage<RoleClient>,
|
||||
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
|
||||
let client = self.client.clone();
|
||||
let message_endpoint = self.message_endpoint.clone();
|
||||
async move {
|
||||
let uri = {
|
||||
let guard = message_endpoint
|
||||
.read()
|
||||
.expect("message endpoint lock poisoned");
|
||||
guard.clone()
|
||||
};
|
||||
client.post_message(uri, item, None).await
|
||||
}
|
||||
}
|
||||
async fn close(&mut self) -> Result<(), Self::Error> {
|
||||
self.stream.take();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: SseClient + std::fmt::Debug> std::fmt::Debug for SseClientTransport<C> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("SseClientWorker")
|
||||
.field("client", &self.client)
|
||||
.field("config", &self.config)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: SseClient> SseClientTransport<C> {
|
||||
pub async fn start_with_client(
|
||||
client: C,
|
||||
config: SseClientConfig,
|
||||
) -> Result<Self, SseTransportError<C::Error>> {
|
||||
let sse_endpoint = config.sse_endpoint.as_ref().parse::<http::Uri>()?;
|
||||
|
||||
let mut sse_stream = client.get_stream(sse_endpoint.clone(), None, None).await?;
|
||||
let initial_message_endpoint = if let Some(endpoint) = config.use_message_endpoint.clone() {
|
||||
let ep = endpoint.parse::<http::Uri>()?;
|
||||
let mut sse_endpoint_parts = sse_endpoint.clone().into_parts();
|
||||
sse_endpoint_parts.path_and_query = ep.into_parts().path_and_query;
|
||||
Uri::from_parts(sse_endpoint_parts)?
|
||||
} else {
|
||||
// wait the endpoint event
|
||||
loop {
|
||||
let sse = sse_stream
|
||||
.next()
|
||||
.await
|
||||
.ok_or(SseTransportError::UnexpectedEndOfStream)??;
|
||||
let Some("endpoint") = sse.event.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
let ep = sse.data.unwrap_or_default();
|
||||
|
||||
break message_endpoint(sse_endpoint.clone(), ep)?;
|
||||
}
|
||||
};
|
||||
let message_endpoint = Arc::new(RwLock::new(initial_message_endpoint));
|
||||
|
||||
let stream = Box::pin(SseAutoReconnectStream::new(
|
||||
sse_stream,
|
||||
SseClientReconnect {
|
||||
client: client.clone(),
|
||||
uri: sse_endpoint.clone(),
|
||||
message_endpoint: message_endpoint.clone(),
|
||||
},
|
||||
config.retry_policy.clone(),
|
||||
));
|
||||
Ok(Self {
|
||||
client,
|
||||
config,
|
||||
message_endpoint,
|
||||
stream: Some(stream),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn message_endpoint(base: http::Uri, endpoint: String) -> Result<http::Uri, http::uri::InvalidUri> {
|
||||
// If endpoint is a full URL, parse and return it directly
|
||||
if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
|
||||
return endpoint.parse::<http::Uri>();
|
||||
}
|
||||
|
||||
let mut base_parts = base.into_parts();
|
||||
let endpoint_clone = endpoint.clone();
|
||||
|
||||
if endpoint.starts_with("?") {
|
||||
// Query only - keep base path and append query
|
||||
if let Some(base_path_and_query) = &base_parts.path_and_query {
|
||||
let base_path = base_path_and_query.path();
|
||||
base_parts.path_and_query = Some(format!("{}{}", base_path, endpoint).parse()?);
|
||||
} else {
|
||||
base_parts.path_and_query = Some(format!("/{}", endpoint).parse()?);
|
||||
}
|
||||
} else {
|
||||
// Path (with optional query) - replace entire path_and_query
|
||||
let path_to_use = if endpoint.starts_with("/") {
|
||||
endpoint // Use absolute path as-is
|
||||
} else {
|
||||
format!("/{}", endpoint) // Make relative path absolute
|
||||
};
|
||||
base_parts.path_and_query = Some(path_to_use.parse()?);
|
||||
}
|
||||
|
||||
http::Uri::from_parts(base_parts).map_err(|_| endpoint_clone.parse::<http::Uri>().unwrap_err())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SseClientConfig {
|
||||
/// client sse endpoint
|
||||
///
|
||||
/// # How this client resolve the message endpoint
|
||||
/// if sse_endpoint has this format: `<schema><authority?><sse_pq>`,
|
||||
/// then the message endpoint will be `<schema><authority?><message_pq>`.
|
||||
///
|
||||
/// For example, if you config the sse_endpoint as `http://example.com/some_path/sse`,
|
||||
/// and the server send the message endpoint event as `message?session_id=123`,
|
||||
/// then the message endpoint will be `http://example.com/message`.
|
||||
///
|
||||
/// This follows the rules of JavaScript's [`new URL(url, base)`](https://developer.mozilla.org/en-US/docs/Web/API/URL/URL)
|
||||
pub sse_endpoint: Arc<str>,
|
||||
pub retry_policy: Arc<dyn SseRetryPolicy>,
|
||||
/// if this is settled, the client will use this endpoint to send message and skip get the endpoint event
|
||||
pub use_message_endpoint: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for SseClientConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
sse_endpoint: "".into(),
|
||||
retry_policy: Arc::new(super::common::client_side_sse::FixedInterval::default()),
|
||||
use_message_endpoint: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures::StreamExt;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct DummyClient;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("dummy error")]
|
||||
struct DummyError;
|
||||
|
||||
impl SseClient for DummyClient {
|
||||
type Error = DummyError;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
_uri: Uri,
|
||||
_message: ClientJsonRpcMessage,
|
||||
_auth_token: Option<String>,
|
||||
) -> Result<(), SseTransportError<Self::Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
_uri: Uri,
|
||||
_last_event_id: Option<String>,
|
||||
_auth_token: Option<String>,
|
||||
) -> Result<BoxedSseResponse, SseTransportError<Self::Error>> {
|
||||
unreachable!("get_stream should not be called in this test")
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_endpoint() {
|
||||
let base_url = "https://localhost/sse".parse::<http::Uri>().unwrap();
|
||||
|
||||
// Query only
|
||||
let result = message_endpoint(base_url.clone(), "?sessionId=x".to_string()).unwrap();
|
||||
assert_eq!(result.to_string(), "https://localhost/sse?sessionId=x");
|
||||
|
||||
// Relative path with query
|
||||
let result = message_endpoint(base_url.clone(), "mypath?sessionId=x".to_string()).unwrap();
|
||||
assert_eq!(result.to_string(), "https://localhost/mypath?sessionId=x");
|
||||
|
||||
// Absolute path with query
|
||||
let result = message_endpoint(base_url.clone(), "/xxx?sessionId=x".to_string()).unwrap();
|
||||
assert_eq!(result.to_string(), "https://localhost/xxx?sessionId=x");
|
||||
|
||||
// Full URL
|
||||
let result = message_endpoint(
|
||||
base_url.clone(),
|
||||
"http://example.com/xxx?sessionId=x".to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(result.to_string(), "http://example.com/xxx?sessionId=x");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_endpoint_control_event_updates_uri() {
|
||||
let initial_endpoint = "https://example.com/message?sessionId=old"
|
||||
.parse::<Uri>()
|
||||
.unwrap();
|
||||
let shared_endpoint = Arc::new(RwLock::new(initial_endpoint));
|
||||
let mut reconnect = SseClientReconnect {
|
||||
client: DummyClient,
|
||||
uri: "https://example.com/sse".parse::<Uri>().unwrap(),
|
||||
message_endpoint: shared_endpoint.clone(),
|
||||
};
|
||||
|
||||
let control_event = Sse::default()
|
||||
.event("endpoint")
|
||||
.data("/message?sessionId=new");
|
||||
|
||||
reconnect.handle_control_event(&control_event).unwrap();
|
||||
|
||||
let guard = shared_endpoint.read().expect("lock poisoned");
|
||||
assert_eq!(
|
||||
guard.to_string(),
|
||||
"https://example.com/message?sessionId=new"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn control_event_frames_are_skipped() {
|
||||
let payload = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"result": {"ok": true}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let events = vec![
|
||||
Ok(Sse::default()
|
||||
.event("endpoint")
|
||||
.data("/message?sessionId=reconnect")),
|
||||
Ok(Sse::default().event("message").data(payload.clone())),
|
||||
];
|
||||
|
||||
let sse_src: BoxedSseResponse = futures::stream::iter(events).boxed();
|
||||
let reconn_stream = SseAutoReconnectStream::never_reconnect(sse_src, DummyError);
|
||||
futures::pin_mut!(reconn_stream);
|
||||
|
||||
let message = reconn_stream.next().await.expect("stream item").unwrap();
|
||||
let actual: Value = serde_json::to_value(message).expect("serialize actual message");
|
||||
// We only need to assert that a valid JSON-RPC response came through after
|
||||
// skipping control frames. The exact `result` shape depends on the SDK's
|
||||
// typed result enums and is not asserted here.
|
||||
assert_eq!(actual.get("jsonrpc"), Some(&Value::String("2.0".into())));
|
||||
assert_eq!(actual.get("id"), Some(&Value::Number(1u64.into())));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,343 +0,0 @@
|
|||
use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
|
||||
|
||||
use axum::{
|
||||
Extension, Json, Router,
|
||||
extract::{NestedPath, Query, State},
|
||||
http::{StatusCode, request::Parts},
|
||||
response::{
|
||||
Response,
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
},
|
||||
routing::{get, post},
|
||||
};
|
||||
use futures::{Sink, SinkExt, Stream};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tokio_util::sync::{CancellationToken, PollSender};
|
||||
use tracing::Instrument;
|
||||
|
||||
use crate::{
|
||||
RoleServer, Service,
|
||||
model::ClientJsonRpcMessage,
|
||||
service::{RxJsonRpcMessage, TxJsonRpcMessage, serve_directly_with_ct},
|
||||
transport::common::server_side_http::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
|
||||
};
|
||||
|
||||
type TxStore =
|
||||
Arc<tokio::sync::RwLock<HashMap<SessionId, tokio::sync::mpsc::Sender<ClientJsonRpcMessage>>>>;
|
||||
pub type TransportReceiver = ReceiverStream<RxJsonRpcMessage<RoleServer>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct App {
|
||||
txs: TxStore,
|
||||
transport_tx: tokio::sync::mpsc::UnboundedSender<SseServerTransport>,
|
||||
post_path: Arc<str>,
|
||||
sse_ping_interval: Duration,
|
||||
}
|
||||
|
||||
impl App {
|
||||
pub fn new(
|
||||
post_path: String,
|
||||
sse_ping_interval: Duration,
|
||||
) -> (
|
||||
Self,
|
||||
tokio::sync::mpsc::UnboundedReceiver<SseServerTransport>,
|
||||
) {
|
||||
let (transport_tx, transport_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
(
|
||||
Self {
|
||||
txs: Default::default(),
|
||||
transport_tx,
|
||||
post_path: post_path.into(),
|
||||
sse_ping_interval,
|
||||
},
|
||||
transport_rx,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PostEventQuery {
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
async fn post_event_handler(
|
||||
State(app): State<App>,
|
||||
Query(PostEventQuery { session_id }): Query<PostEventQuery>,
|
||||
parts: Parts,
|
||||
Json(mut message): Json<ClientJsonRpcMessage>,
|
||||
) -> Result<StatusCode, StatusCode> {
|
||||
tracing::debug!(session_id, ?parts, ?message, "new client message");
|
||||
let tx = {
|
||||
let rg = app.txs.read().await;
|
||||
rg.get(session_id.as_str())
|
||||
.ok_or(StatusCode::NOT_FOUND)?
|
||||
.clone()
|
||||
};
|
||||
message.insert_extension(parts);
|
||||
if tx.send(message).await.is_err() {
|
||||
tracing::error!("send message error");
|
||||
return Err(StatusCode::GONE);
|
||||
}
|
||||
Ok(StatusCode::ACCEPTED)
|
||||
}
|
||||
|
||||
async fn sse_handler(
|
||||
State(app): State<App>,
|
||||
nested_path: Option<Extension<NestedPath>>,
|
||||
parts: Parts,
|
||||
) -> Result<Sse<impl Stream<Item = Result<Event, io::Error>>>, Response<String>> {
|
||||
let session = session_id();
|
||||
tracing::info!(%session, ?parts, "sse connection");
|
||||
use tokio_stream::{StreamExt, wrappers::ReceiverStream};
|
||||
use tokio_util::sync::PollSender;
|
||||
let (from_client_tx, from_client_rx) = tokio::sync::mpsc::channel(64);
|
||||
let (to_client_tx, to_client_rx) = tokio::sync::mpsc::channel(64);
|
||||
let to_client_tx_clone = to_client_tx.clone();
|
||||
|
||||
app.txs
|
||||
.write()
|
||||
.await
|
||||
.insert(session.clone(), from_client_tx);
|
||||
let session = session.clone();
|
||||
let stream = ReceiverStream::new(from_client_rx);
|
||||
let sink = PollSender::new(to_client_tx);
|
||||
let transport = SseServerTransport {
|
||||
stream,
|
||||
sink,
|
||||
session_id: session.clone(),
|
||||
tx_store: app.txs.clone(),
|
||||
};
|
||||
let transport_send_result = app.transport_tx.send(transport);
|
||||
if transport_send_result.is_err() {
|
||||
tracing::warn!("send transport out error");
|
||||
let mut response =
|
||||
Response::new("fail to send out transport, it seems server is closed".to_string());
|
||||
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
|
||||
return Err(response);
|
||||
}
|
||||
let nested_path = nested_path.as_deref().map(NestedPath::as_str).unwrap_or("");
|
||||
let post_path = app.post_path.as_ref();
|
||||
let ping_interval = app.sse_ping_interval;
|
||||
let stream = futures::stream::once(futures::future::ok(
|
||||
Event::default()
|
||||
.event("endpoint")
|
||||
.data(format!("{nested_path}{post_path}?sessionId={session}")),
|
||||
))
|
||||
.chain(ReceiverStream::new(to_client_rx).map(|message| {
|
||||
match serde_json::to_string(&message) {
|
||||
Ok(bytes) => Ok(Event::default().event("message").data(&bytes)),
|
||||
Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
|
||||
}
|
||||
}));
|
||||
|
||||
tokio::spawn(async move {
|
||||
// Wait for connection closure
|
||||
to_client_tx_clone.closed().await;
|
||||
|
||||
// Clean up session
|
||||
let session_id = session.clone();
|
||||
let tx_store = app.txs.clone();
|
||||
let mut txs = tx_store.write().await;
|
||||
txs.remove(&session_id);
|
||||
tracing::debug!(%session_id, "Closed session and cleaned up resources");
|
||||
});
|
||||
|
||||
Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(ping_interval)))
|
||||
}
|
||||
|
||||
pub struct SseServerTransport {
|
||||
stream: ReceiverStream<RxJsonRpcMessage<RoleServer>>,
|
||||
sink: PollSender<TxJsonRpcMessage<RoleServer>>,
|
||||
session_id: SessionId,
|
||||
tx_store: TxStore,
|
||||
}
|
||||
|
||||
impl Sink<TxJsonRpcMessage<RoleServer>> for SseServerTransport {
|
||||
type Error = io::Error;
|
||||
|
||||
fn poll_ready(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Result<(), Self::Error>> {
|
||||
self.sink
|
||||
.poll_ready_unpin(cx)
|
||||
.map_err(std::io::Error::other)
|
||||
}
|
||||
|
||||
fn start_send(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
item: TxJsonRpcMessage<RoleServer>,
|
||||
) -> Result<(), Self::Error> {
|
||||
self.sink
|
||||
.start_send_unpin(item)
|
||||
.map_err(std::io::Error::other)
|
||||
}
|
||||
|
||||
fn poll_flush(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Result<(), Self::Error>> {
|
||||
self.sink
|
||||
.poll_flush_unpin(cx)
|
||||
.map_err(std::io::Error::other)
|
||||
}
|
||||
|
||||
fn poll_close(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Result<(), Self::Error>> {
|
||||
let inner_close_result = self
|
||||
.sink
|
||||
.poll_close_unpin(cx)
|
||||
.map_err(std::io::Error::other);
|
||||
if inner_close_result.is_ready() {
|
||||
let session_id = self.session_id.clone();
|
||||
let tx_store = self.tx_store.clone();
|
||||
tokio::spawn(async move {
|
||||
tx_store.write().await.remove(&session_id);
|
||||
});
|
||||
}
|
||||
inner_close_result
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for SseServerTransport {
|
||||
type Item = RxJsonRpcMessage<RoleServer>;
|
||||
|
||||
fn poll_next(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Self::Item>> {
|
||||
use futures::StreamExt;
|
||||
self.stream.poll_next_unpin(cx)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SseServerConfig {
|
||||
pub bind: SocketAddr,
|
||||
pub sse_path: String,
|
||||
pub post_path: String,
|
||||
pub ct: CancellationToken,
|
||||
pub sse_keep_alive: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct SseServer {
|
||||
transport_rx: tokio::sync::mpsc::UnboundedReceiver<SseServerTransport>,
|
||||
pub config: SseServerConfig,
|
||||
}
|
||||
|
||||
impl SseServer {
|
||||
pub async fn serve(bind: SocketAddr) -> io::Result<Self> {
|
||||
Self::serve_with_config(SseServerConfig {
|
||||
bind,
|
||||
sse_path: "/sse".to_string(),
|
||||
post_path: "/message".to_string(),
|
||||
ct: CancellationToken::new(),
|
||||
sse_keep_alive: None,
|
||||
})
|
||||
.await
|
||||
}
|
||||
pub async fn serve_with_config(config: SseServerConfig) -> io::Result<Self> {
|
||||
let (sse_server, service) = Self::new(config);
|
||||
let listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
|
||||
let ct = sse_server.config.ct.child_token();
|
||||
let server = axum::serve(listener, service).with_graceful_shutdown(async move {
|
||||
ct.cancelled().await;
|
||||
tracing::info!("sse server cancelled");
|
||||
});
|
||||
tokio::spawn(
|
||||
async move {
|
||||
if let Err(e) = server.await {
|
||||
tracing::error!(error = %e, "sse server shutdown with error");
|
||||
}
|
||||
}
|
||||
.instrument(tracing::info_span!("sse-server", bind_address = %sse_server.config.bind)),
|
||||
);
|
||||
Ok(sse_server)
|
||||
}
|
||||
|
||||
pub fn new(config: SseServerConfig) -> (SseServer, Router) {
|
||||
let (app, transport_rx) = App::new(
|
||||
config.post_path.clone(),
|
||||
config.sse_keep_alive.unwrap_or(DEFAULT_AUTO_PING_INTERVAL),
|
||||
);
|
||||
let router = Router::new()
|
||||
.route(&config.sse_path, get(sse_handler))
|
||||
.route(&config.post_path, post(post_event_handler))
|
||||
.with_state(app);
|
||||
|
||||
let server = SseServer {
|
||||
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(std::io::Error::other)?;
|
||||
server.waiting().await?;
|
||||
tokio::io::Result::Ok(())
|
||||
});
|
||||
}
|
||||
});
|
||||
ct
|
||||
}
|
||||
|
||||
/// This allows you to skip the initialization steps for incoming request.
|
||||
pub fn with_service_directly<S, F>(mut self, service_provider: F) -> CancellationToken
|
||||
where
|
||||
S: Service<RoleServer>,
|
||||
F: Fn() -> S + Send + 'static,
|
||||
{
|
||||
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 = serve_directly_with_ct(service, transport, None, ct);
|
||||
server.waiting().await?;
|
||||
tokio::io::Result::Ok(())
|
||||
});
|
||||
}
|
||||
});
|
||||
ct
|
||||
}
|
||||
|
||||
pub fn cancel(&self) {
|
||||
self.config.ct.cancel();
|
||||
}
|
||||
|
||||
pub async fn next_transport(&mut self) -> Option<SseServerTransport> {
|
||||
self.transport_rx.recv().await
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for SseServer {
|
||||
type Item = SseServerTransport;
|
||||
|
||||
fn poll_next(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Self::Item>> {
|
||||
self.transport_rx.poll_recv(cx)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue