modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

refactor(sse_server): separate router and server startup (#52)

Browse source 
* feat: expose axum router

* feat: add axum_router example with SSE server implementation

* refactor: simplify SseServer configuration handling in server setup

* docs: add warning to SseServer::new about potential post_path issues with embedded routers
This commit is contained in:
Wang Ruobiao 2025-03-28 22:06:50 +08:00 • committed by GitHub
commit 59bfb6e748
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 77 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

31
crates/rmcp/src/transport/sse_server.rs
View file

@ -219,13 +219,9 @@ impl SseServer {
.await
}
pub async fn serve_with_config(config: SseServerConfig) -> io::Result<Self> {
let (app, transport_rx) = App::new(config.post_path.clone());
let listener = tokio::net::TcpListener::bind(config.bind).await?;
let service = Router::new()
.route(&config.sse_path, get(sse_handler))
.route(&config.post_path, post(post_event_handler))
.with_state(app);
let ct = config.ct.child_token();
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");
@ -236,13 +232,28 @@ impl SseServer {
tracing::error!(error = %e, "sse server shutdown with error");
}
}
.instrument(tracing::info_span!("sse-server", bind_address = %config.bind)),
.instrument(tracing::info_span!("sse-server", bind_address = %sse_server.config.bind)),
);
Ok(Self {
Ok(sse_server)
}
/// Warning: This function creates a new SseServer instance with the provided configuration.
/// `App.post_path` may be incorrect if using `Router` as an embedded router.
pub fn new(config: SseServerConfig) -> (SseServer, Router) {
let (app, transport_rx) = App::new(config.post_path.clone());
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>,